Method _register_python_functions

python/tvm/relax/base_py_module.py:192–219

Register Python functions with the VM runtime for call_py_func support.

(self)

            setattr(self, func_name, _create_relax_wrapper(func_name))

    def _register_python_functions(self):
        """Register Python functions with the VM runtime for call_py_func support."""
        if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs:
            return

        try:
            register_py_func = tvm.get_global_func("vm.builtin.register_py_func")
        except ValueError:
            return

        for func_name, py_func in self.ir_mod.pyfuncs.items():

            def create_py_func_wrapper(name, original_func):
                def wrapper(*args, **kwargs):
                    converted_args = [self._convert_tvm_to_pytorch(arg) for arg in args]
                    converted_kwargs = {
                        k: self._convert_tvm_to_pytorch(v) for k, v in kwargs.items()
                    }

                    result = original_func(self, *converted_args, **converted_kwargs)

                    return self._convert_pytorch_to_tvm(result)

                wrapper.__name__ = name
                return wrapper

            wrapped_func = create_py_func_wrapper(func_name, py_func)
            register_py_func(func_name, wrapped_func)

    def call_tir(self, tir_func, args, out_sinfo):
        """Call a TIR function with PyTorch tensors."""