Function test_where

tests/python/relax/test_tvmscript_parser_op_search.py:37–55

()

def test_where():
    @R.function
    def foo(
        condition: R.Tensor((2, 1), "bool"),
        x: R.Tensor((2, 3), "float32"),
        y: R.Tensor((1, 3), "float32"),
    ) -> R.Tensor((2, 3), "float32"):
        gv: R.Tensor((2, 3), "float32") = R.where(condition, x, y)
        return gv

    bb = relax.BlockBuilder()
    condition = relax.Var("condition", R.Tensor((2, 1), "bool"))
    x = relax.Var("x", R.Tensor((2, 3), "float32"))
    y = relax.Var("y", R.Tensor((1, 3), "float32"))
    with bb.function("foo", [condition, x, y]):
        gv = bb.emit(relax.op.where(condition, x, y))
        bb.emit_func_output(gv)

    _check(foo, bb.get()["foo"])


def test_argmax():