Class ShapeExpr

python/tvm/relax/expr.py:684–712

A shape expression which allows users to construct a shape containing PrimExpr. Parameters ---------- values: Union[List[PrimExpr], typing.Tuple[PrimExpr, ...], tvm_ffi.Array] The values of the shape expression. span: Optional[Span] Span that points to original sour

@tvm_ffi.register_object("relax.expr.ShapeExpr")
class ShapeExpr(ExprWithOp):
    """A shape expression which allows users to construct a shape containing PrimExpr.

    Parameters
    ----------
    values: Union[List[PrimExpr], typing.Tuple[PrimExpr, ...], tvm_ffi.Array]
        The values of the shape expression.

    span: Optional[Span]
        Span that points to original source code
    """

    values: list[PrimExpr]
    span: Span | None

    def __init__(
        self,
        values: list[PrimExpr] | tuple[PrimExpr, ...] | tvm_ffi.Array,
        span: Span | None = None,
    ) -> None:
        self.__init_handle_by_constructor__(_ffi_api.ShapeExpr, values, span)  # type: ignore

    def __getitem__(self, index):
        if index >= len(self) or index < -len(self):
            raise IndexError("ShapeExpr index out of range")
        return self.values[index]

    def __len__(self):
        return len(self.values)


def make_shape(shape: list[Any] | tuple[Any, ...]) -> ShapeExpr: