Skip to main content

Mojo 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, 0]): 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, 0])

Constructs a ranked tensor type.

Args:

  • dtype (DType): The element type of the tensor data.
  • dims (List[Dim, 0]): 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.