Function wrap_nested

python/tvm/relax/frontend/nn/core.py:762–791

Wrap the given relax.Expr, emit it using the current BlockBuilder, and automatically handle nested cases if the expr represents a Tuple. Parameters ---------- expr : relax.Expr The Expr to be wrapped. name : str Name hint. Returns ------- result : U

(expr: rx.Expr, name: str)

def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
    """Wrap the given relax.Expr, emit it using the current BlockBuilder,
    and automatically handle nested cases if the expr represents a Tuple.

    Parameters
    ----------
    expr : relax.Expr
        The Expr to be wrapped.

    name : str
        Name hint.

    Returns
    -------
    result : Union[Tensor, Tuple[Tensor]]
        The computed result.
    """
    if not isinstance(expr, rx.DataflowVar):
        expr = BlockBuilder.current().emit(expr, name)
    if isinstance(expr.struct_info_, TensorStructInfo):
        return Tensor(_expr=expr)
    if isinstance(expr.struct_info_, TupleStructInfo):
        return tuple(
            wrap_nested(  # type: ignore
                rx.TupleGetItem(expr, i),
                name=f"{name}.{i}",
            )
            for i in range(len(expr.struct_info_.fields))
        )
    raise TypeError(f"Unsupported return type: {expr.struct_info_}")


def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any], bool]):