struct
TensorType
A symbolic tensor type.
It is not an eager tensor type!! It contains no actual data, but instead represents a value at some point in time during model execution.
Most internal values in a model will be tensors. This type represents their element type (dtype) and dimensions (dims) at a specific point during model computation. It allows us to do some optimistic optimizations and shape inference during graph construction, and to provide more detailed shape information to the compiler for further optimizatino passes.
It can also represent a fully dynamic rank tensor. The presence of dynamic rank tensors in a graph will often degrade performance dramatically and prevents many classes of optimizations.
Fields
- dtype (
DType): The element type of the tensor value. - dims (
List[Dim]): The dimensions of the tensor value, if it is known-rank.
Implemented traits
AnyType,
CollectionElement,
Copyable,
Movable
Methods
__init__
__init__(inout self: Self, dtype: DType)
Constructs a 0-d tensor type.
Args:
- dtype (
DType): The element type of the tensor data.
__init__(inout self: Self, dtype: DType, *dims: Dim)
Constructs a tensor type.
Args:
- dtype (
DType): The element type of the tensor data. - *dims (
Dim): The shape dimensions of the tensor. The number of dims is the rank of the tensor.
__init__(inout self: Self, dtype: DType, dims: List[Dim])
Constructs a ranked tensor type.
Args:
- dtype (
DType): The element type of the tensor data. - dims (
List[Dim]): The shape dimensions of the tensor. The number of dims is the rank of the tensor.
__init__(inout self: Self, spec: TensorSpec)
Constructs a tensor type from a TensorSpec.
Since TensorSpec can only contain static shapes, this will always construct a static tensor.
Args:
- spec (
TensorSpec): The dtype and static shape of the tensor.
__eq__
__eq__(self: Self, other: Self) -> Bool
Checks whether the two tensors are identical (same rank, type, shape).
Args:
- other (
Self): The other tensor to check equality against.
Returns:
True if the tensors have identical element type and shape, False otherwise.
to_mlir
to_mlir(self: Self, ctx: Context) -> Type
Converts to an _mlir.Type instance.
Args:
- ctx (
Context): The mlir.Context in which to create the type.
Returns:
An _mlir.Type in the specified Context.
from_mlir
static from_mlir(t: Type) -> Self
Constructs a tensor type from an _mlir type.
Args:
- t (
Type): The _mlir Type object to parse into a tensor type.
Returns:
The tensor type represented by the _mlir Type value.
is_static
is_static(self: Self) -> Bool
Checks whether the tensor type has a fully static shape or not.
A tensor must have all of its dimensions be static (or be 0-dimensional)
in order to be static.
Returns:
True if the tensor has a fully static shape, False otherwise.
rank
rank(self: Self) -> Int
Gets the rank of the tensor type.
Returns:
The tensor's static rank, or 0 for a dynamic tensor. A 0-dimensional static tensor also has a rank of 0, so check ranked directly to check if a tensor is ranked or not.
dim
dim(self: Self, pos: Int) -> Dim
Gets the pos'th dimension of the tensor type.
Supports negative-indexing, ie. t.dim(-1) will give the last
dimension.
Args:
- pos (
Int): The dimension index to retrieve.
Returns:
The dimension value at dimension pos.
Raises:
If the dimension is out-of-bounds, or if the tensor is unranked.
num_elements
num_elements(self: Self) -> SIMD[int64, 1]
Counts the total number of elements in the tensor type.
For a static tensor, returns the product of all static dimensions. This is the number of elements the tensor will hold during execution, TensorType doesn't actually hold any element values at all.
For any non-static tensor, ie. a tensor having any symbolic or dynamic dimensions, the return value will be meaningless.
Returns:
The number of elements the tensor contains.
cast
cast(self: Self, dtype: DType) -> Self
Constructs a new tensor type of the same shape with the new dtype.
Args:
- dtype (
DType): The new element type for the tensor.
Returns:
A new tensor type with the same shape, and the new element type.