Function test_call_packed

tests/python/relax/test_ast_printer.py:349–416

()

def test_call_packed():
    # test case from test_parser
    @R.function(pure=False)
    def f(
        x: R.Tensor((32, "m"), "float32"),
        y: R.Tensor(("m",), "float32"),
        r: R.Tensor(dtype="int64"),
    ) -> R.Object:
        m = T.int64()
        z: R.Tensor((32, m), "float32") = R.multiply(x, y)
        w: R.Tensor(ndim=2) = R.multiply(z, z)
        q: R.Tensor = R.add(w, w)
        t = R.add(w, z)
        sh: R.Shape = R.shape_of(t)
        o: R.Object = R.call_packed(
            "contrib.tensor_array_stack", x, y, sinfo_args=R.Object(), test_attr=True
        )
        return o

    # checking that the call_packed call is turned into a call to an extern func
    f_str = strip_whitespace(
        dump_ast(
            f,
            include_struct_info_annotations=False,
            include_call_attrs=True,
        )
    )

    # the function has an annotated return type
    assert "ret_struct_info=ObjectStructInfo()" in f_str
    # the purity attribute is set to false
    assert "is_pure=False"

    assert isinstance(f.body, rx.SeqExpr)
    extern_call = f.body.blocks[0].bindings[-1].value
    extern_call_text = dump_ast(
        extern_call,
        include_struct_info_annotations=False,
        include_call_attrs=True,
    )
    assert strip_whitespace(extern_call_text) in f_str
    assert_fields(
        "Call",
        {
            "op": 'ExternFunc(global_symbol="contrib.tensor_array_stack")',
            "args": '[Var(name_hint="x"), Var(name_hint="y")]',
            "sinfo_args": "[ObjectStructInfo()]",
            "attrs": '{"test_attr": True}',
        },
        extern_call_text,
    )

    # check that the op call is there too
    op_call = f.body.blocks[0].bindings[0].value
    op_call_text = dump_ast(
        op_call,
        include_struct_info_annotations=False,
        include_call_attrs=True,