Function call_py_func

python/tvm/relax/script/builder/ir.py:464–512

Create a relax Call, which calls a Python function. Parameters ---------- py_func_name: str The name of the Python function to call. This should correspond to a function in the IRModule's pyfuncs attribute. *args : Expr The arguments. out_sinfo: Union[Str

(
    py_func_name: py_str,
    *args: Expr,
    out_sinfo: StructInfo | list[StructInfo],
)

def call_py_func(
    py_func_name: py_str,
    *args: Expr,
    out_sinfo: StructInfo | list[StructInfo],
) -> Call:
    """Create a relax Call, which calls a Python function.

    Parameters
    ----------
    py_func_name: str
        The name of the Python function to call. This should correspond to a function
        in the IRModule's pyfuncs attribute.
    *args : Expr
        The arguments.
    out_sinfo: Union[StructInfo, List[StructInfo]]
        The structure info of the call_py_func output.
        It should be a single or a list of TensorStructInfo. Each one denotes the
        structure info of a returned tensor.

    Returns
    -------
    call: Call
        The created Relax Call for call_py_func operator.
    """
    args = py_tuple(convert_to_expr(a) for a in args)
    if isinstance(out_sinfo, py_tuple):  # type: ignore
        out_sinfo = list(out_sinfo)
    elif not isinstance(out_sinfo, list):
        out_sinfo = [out_sinfo]

    out_sinfo = [
        (
            sinfo()
            if callable(sinfo)
            else sinfo.asobject()
            if isinstance(sinfo, ObjectConvertible)
            else sinfo
        )
        for sinfo in out_sinfo
    ]

    # Convert string to StringImm
    try:
        func_name_imm = (
            StringImm(py_func_name) if isinstance(py_func_name, py_str) else py_func_name
        )
    except (TypeError, ValueError, AttributeError):
        func_name_imm = StringImm(py_func_name)
    return _call_py_func(func_name_imm, args, out_sinfo)


def _sinfo_arg_wrapper(func):