Class TensorStructInfo

python/tvm/relax/struct_info.py:142–187

StructInfo of a Tensor value. Parameters ---------- shape : Optional[Expr] The shape expression. dtype : Optional[str] The content data type. vdevice : Optional[Vdevice] The virtual device. ndim : Optional[int] The number of dimensions of the

@tvm_ffi.register_object("relax.TensorStructInfo")
class TensorStructInfo(StructInfo):
    """StructInfo of a Tensor value.

    Parameters
    ----------
    shape : Optional[Expr]
       The shape expression.

    dtype : Optional[str]
        The content data type.

    vdevice : Optional[Vdevice]
        The virtual device.

    ndim : Optional[int]
       The number of dimensions of the tensor.

    Note
    ----
    Do not specify shape and ndim at the same time.
    """

    shape: Expr | None
    dtype: str
    vdevice: VDevice | None
    ndim: int
    span: Span

    def __init__(
        self,
        shape: Expr | None | list[PrimExpr] = None,
        dtype: str = "float32",
        vdevice: VDevice | None | str = None,
        ndim: int = -1,
        span: Span = None,
    ) -> None:
        if isinstance(shape, list | tuple | Array):
            shape = ShapeExpr(shape)
        self.__init_handle_by_constructor__(
            _ffi_api.TensorStructInfo,
            shape,
            dtype,
            ndim,
            vdevice,
            span,  # type: ignore
        )


@tvm_ffi.register_object("relax.TupleStructInfo")