This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvmscript-generic-parser-builder in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c2f4133b011fe02af166b9795417ce547c05db9c Author: Tianqi Chen <[email protected]> AuthorDate: Mon Sep 21 02:33:44 2026 +0000 Register the concrete Python module factory at the Relax entry --- python/tvm/relax/script/v2.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/tvm/relax/script/v2.py b/python/tvm/relax/script/v2.py index 9030ca0d2d..1b627f04b5 100644 --- a/python/tvm/relax/script/v2.py +++ b/python/tvm/relax/script/v2.py @@ -20,6 +20,7 @@ import sys as _sys from tvm import relax as _relax +from tvm.relax.base_py_module import PyModuleFactory as _PyModuleFactory from tvm.script.parser_v2.frontend import make_decorator as _make_decorator from tvm.script.parser_v2.frontend import make_helper as _make_helper from tvm.script.parser_v2.frontend import register_namespace as _register_namespace @@ -47,4 +48,11 @@ def _opaque_function(name, function, source, span): ) -_register_opaque_factory(_opaque_function) +def _python_module(module, original, bases): + candidates = (original, *bases) + if any(isinstance(base, type) and issubclass(base, _relax.BasePyModule) for base in candidates): + return _PyModuleFactory(module, original_class=original) + return module + + +_register_opaque_factory(_opaque_function, module_adapter=_python_module)
