Class ModuleFactory

python/tvm/script/parser/ir/entry.py:108–138

Factory class for creating BasePyModule instances with Python functions.

                    continue

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

                def __init__(self, module, pyfunc_methods, original_class):
                    self.ir_module = module
                    self.pyfunc_methods = pyfunc_methods
                    self.original_class = original_class

                def __call__(self, device=None, target=None):
                    if device is None:
                        device = cpu(0)

                    instance_ir_mod = ir.IRModule()
                    for global_var, func in self.ir_module.functions_items():
                        instance_ir_mod[global_var] = func

                    instance = BasePyModule(instance_ir_mod, device, target)

                    for method_name in self.pyfunc_methods:
                        if hasattr(self.original_class, method_name):
                            method = getattr(self.original_class, method_name)
                            instance.add_python_function(method_name, method)

                    return instance

                def __getattr__(self, name):
                    if hasattr(self.ir_module, name):
                        return getattr(self.ir_module, name)
                    raise AttributeError(
                        f"'{self.__class__.__name__}' object has no attribute '{name}'"
                    )

            factory = ModuleFactory(m, pyfunc_methods, mod)
            setattr(factory, "__name__", mod.__name__)