Function decorator_wrapper

python/tvm/script/parser/ir/entry.py:54–145

(mod)

    outer_stack = inspect.stack()

    def decorator_wrapper(mod):
        if not inspect.isclass(mod):
            raise TypeError(f"Expect a class, but got: {mod}")

        # Check BasePyModule inheritance
        base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)

        extra_vars = utils.inspect_class_capture(mod)
        # Resolve closure variables hidden by PEP 563 (annotation-only names)
        utils.resolve_closure_vars(mod, extra_vars, outer_stack)
        m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)

        if base_py_module_inherited:
            # Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
            # because tvm.script is loaded before tvm.relax during tvm initialization.
            from tvm.relax.base_py_module import BasePyModule
            from tvm.relax.expr import ExternFunc  # pylint: disable=import-outside-toplevel

            # Collect pyfunc methods
            pyfunc_methods = [
                name
                for name, attr in mod.__dict__.items()
                if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
            ]

            mod._pyfunc_methods = pyfunc_methods

            # Create ExternFunc nodes

            for method_name in pyfunc_methods:
                try:
                    existing_gvars = [
                        global_var
                        for global_var in m.get_global_vars()
                        if global_var.name_hint == method_name
                    ]

                    extern_func = ExternFunc(method_name)
                    extern_func = extern_func.with_attr("is_pyfunc", True)
                    extern_func = extern_func.with_attr("function_type", "python")
                    extern_func = extern_func.with_attr("python_function_name", method_name)
                    extern_func = extern_func.with_attr(
                        "python_source", f"# Source for {method_name}"
                    )
                    extern_func = extern_func.with_attr("python_packed_func", None)

                    if existing_gvars:
                        m[existing_gvars[0]] = extern_func
                    else:
                        m[GlobalVar(method_name)] = extern_func

                except Exception:  # pylint: disable=broad-exception-caught
                    continue

            class ModuleFactory:
                """Factory class for creating BasePyModule instances with Python functions."""

                def __init__(self, module, pyfunc_methods, original_class):