Skip to content

Support ast.IfExp in kernels #500

@shi-eric

Description

@shi-eric

Kernels should support expressions like out[i] = 1.0 if out[i] < value else -1, e.g.

import warp as wp

wp.init()


@wp.kernel
def threshold(value: float, out: wp.array(dtype=float)):
    i = wp.tid()

    out[i] = 1.0 if out[i] < value else -1


a = wp.ones(10, dtype=float)
wp.launch(threshold, (10,), inputs=[0.5, a])

print(a.numpy())

Currently, the above code yields:

warp.codegen.WarpCodegenError: Error while parsing function "threshold" at /home/eshi/code-projects/warp/ershi-tests/test_single_line.py:10:
    out[i] = 1.0 if out[i] < value else -1
;Construct `ast.IfExp` not supported in kernels.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions