Function test_extern_func

tests/python/relax/test_tvmscript_printer_relax.py:75–99

()

def test_extern_func():
    @R.function
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):  # type: ignore
        return a

    obj = IRModule(
        {
            "func": func,
            "my_ext": relax.ExternFunc("my_ext"),
        }
    )
    _assert_print(
        obj,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    my_ext = R.ExternFunc("my_ext")
    @R.function
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
        return a
""",
    )


def test_extern_func_with_struct_info():