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
@Hzfengsy Hzfengsy deleted the tvmscript_ifexp branch October 11, 2025 03:49
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