Function _attach_pyfuncs_to_irmodule

python/tvm/script/parser/core/entry.py:180–210

Attach Python functions to IRModule with reduced nesting.

(irmodule, all_pyfuncs)

def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
    """Attach Python functions to IRModule with reduced nesting."""
    if not all_pyfuncs:
        return

    if not hasattr(irmodule, "pyfuncs"):
        irmodule.pyfuncs = {}

    for global_var, func in irmodule.functions_items():
        if not isinstance(func, tvm.relax.ExternFunc):
            continue
        if not func.attrs.get("is_pyfunc", False):
            continue

        pyfunc_name = global_var.name_hint
        if pyfunc_name not in all_pyfuncs:
            continue

        pyfunc = all_pyfuncs[pyfunc_name]
        irmodule.pyfuncs[pyfunc_name] = pyfunc

        try:
            source_code = inspect.getsource(pyfunc)
            func = func.with_attr("python_source", source_code)
        except (OSError, TypeError):
            func = func.with_attr("python_source", f"# Source unavailable for {pyfunc_name}")

        packed_func = _create_python_packed_func(pyfunc)
        func = func.with_attr("python_packed_func", packed_func)

        irmodule[global_var] = func