Method call_dps_packed

python/tvm/relax/base_py_module.py:256–275

Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack.

(self, func_name: str, args, out_sinfo)

        return result[0] if len(result) == 1 else result

    def call_dps_packed(self, func_name: str, args, out_sinfo):
        """Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack."""
        if hasattr(self, func_name) and callable(getattr(self, func_name)):
            return getattr(self, func_name)(*args)

        if func_name not in self.extern_funcs:
            try:
                self.extern_funcs[func_name] = tvm.get_global_func(func_name)
            except ValueError as error:
                raise ValueError(
                    f"Function '{func_name}' not found as a global function. "
                    f"Please implement it as a method or register it."
                ) from error
        func = self.extern_funcs[func_name]

        out = self._create_output_tensors(out_sinfo, args)
        tvm_args = self._convert_pytorch_to_tvm(args)
        tvm_out = self._convert_pytorch_to_tvm(out)
        func(*tvm_args, *tvm_out)
        return out[0] if len(out) == 1 else out

    def call_py_func(self, func_name: str, args):
        """Call a Python function stored in the module's pyfuncs."""