Skip to main content

/

Docs

Log in

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.

Was this page helpful?