Function test_nested_function

tests/python/relax/test_tvmscript_printer_relax.py:141–172

()

def test_nested_function():
    @I.ir_module(s_tir=True)
    class NestedFunction:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            @R.function
            def nested(y: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
                return y

            z = nested(x)
            return z

    _assert_print(
        NestedFunction,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        # from tvm.script import relax as R

        @R.function
        def nested(y: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
            return y

        z: R.Tensor((), dtype="int32") = nested(x)
        return z
""",
    )


def test_object_struct_info():