Function gen_mod

tests/python/relax/test_transform_fold_constant.py:29–57

Select relax function with name, rename to main and bind constant. Parameters ---------- mod: IRModule The input module name: str The name of relax function to preserve and rename to main binding: Dict[str, array] The const parameter bindings

(mod, name, binding)

def gen_mod(mod, name, binding):
    """Select relax function with name, rename to main and bind constant.

    Parameters
    ----------
    mod: IRModule
        The input module

    name: str
        The name of relax function to preserve and rename to main

    binding: Dict[str, array]
        The const parameter bindings
    """
    funcs = {}
    binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()}

    for k, v in mod.functions.items():
        if isinstance(v, tvm.relax.Function):
            if k.name_hint == name:
                # rename to main
                gv = tvm.ir.GlobalVar("main")
                funcs[gv] = tvm.relax.Function(v.params, v.body, v.ret_struct_info).with_attr(
                    "global_symbol", "main"
                )
        else:
            funcs[k] = v
    mod = tvm.IRModule(funcs)
    return relax.transform.BindParams("main", binding)(mod)


def test_one_fold_addone():