Skip to content

Conversation

Hzfengsy
Copy link
Member

This PR adds support for conditional expressions in TVMScript parser, which allows developers to use Python-style conditional expressions

@T.prim_func
def func(A: T.buffer((128, 128), "float32")):
    for i, j in T.grid(128, 128):
        A[i, j] = i if i < j else j

@T.prim_func
def expected(A: T.buffer((128, 128), "float32")):
    for i, j in T.grid(128, 128):
        A[i, j] = T.if_then_else(i < j, i, j)

@Hzfengsy Hzfengsy changed the title Add support for conditional expressions in TVMScript [TIR] Add support for conditional expressions in TVMScript Sep 19, 2025
This PR adds support for conditional expressions in TVMScript parser,
which allows developers to use Python-style conditional expressions

```python
@T.prim_func
def func(A: T.buffer((128, 128), "float32")):
    for i, j in T.grid(128, 128):
        A[i, j] = i if i < j else j

@T.prim_func
def expected(A: T.buffer((128, 128), "float32")):
    for i, j in T.grid(128, 128):
        A[i, j] = T.if_then_else(i < j, i, j)
```
Copy link
Contributor

@MasterJH5574 MasterJH5574 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for introducing the sugar!

@MasterJH5574 MasterJH5574 merged commit 7ec2d35 into apache:main Sep 20, 2025
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants