Function test_operators

tests/python/relax/test_ast_printer.py:553–581

()

def test_operators():
    @R.function
    def foo(x: R.Tensor):
        return R.unique(x, sorted=True, axis=-1)

    foo_str = strip_whitespace(
        dump_ast(
            foo,
            include_struct_info_annotations=False,
        )
    )
    assert 'Op(name="relax.unique")' in foo_str
    # the sorted argument is true, so it will be a PrimValue of 1
    assert "PrimExpr(value=`T.int64(1)`)" in foo_str
    # axis is -1
    assert "PrimExpr(value=`T.int64(-1)`)" in foo_str

    @R.function(pure=False)
    def bar(x: R.Tensor):
        return R.print(x, format="{}")

    bar_str = strip_whitespace(
        dump_ast(
            bar,
            include_struct_info_annotations=False,
        )
    )
    # the format string is a StringImm argument
    assert 'StringImm(value="{}")' in bar_str


def test_print_struct_info_annotation_non_var():