Skip to main content

type

Library for graph Symbol Types.

DynamicDim

A dynamic tensor dimension.

DynamicDims are printed in MO tensor types as ?, eg. !mo.tensor<[?, 4, ?], si32]> has 2 dynamic dimensions.

Dynamic dimensions reduce the compiler's ability to reason about tensor shapes as data moves through the model, and may therefore limit the available optimizations it can perform. Reducing usage of dynamic dims can be an avenue to improving model performance.

Create a dynamic dimension via Dim.dynamic().

Implemented traits:

AnyType, CollectionElement, Copyable, Movable

SymbolicDim

A symbolic tensor dimension.

SymbolicDimss have a name and are printed as their name on MO types, eg. !mo.tensor<[batch, x, 10], si32]> the first and second dimensions are named "batch" and "x" respectively.

Symbolic dimensions don't have a static value, but they allow a readable name to understand what's going on in the model IR better, and they also allow users to hint to the compiler that two dimensions will have the same value, which can often allow important speedups.

Create a symbolic dimension via Dim.symbolic("name"), or just by passing the string name, eg. MOTensor(DType.bool, "batch", Dim.dynamic(), 10).

Fields:

  • name (String): The name of the dimension.

Implemented traits:

AnyType, CollectionElement, Copyable, Movable

Methods:

__eq__

__eq__(self: Self, other: Self) -> Bool

Whether the dimension is the same as another symbolic dimension.

Symbolic dimensions with the same name are interpreted as the same dimensionality! If you use Symbolic dimensions, make sure you're naming them consistently, your model will likely fail to compile if you name two actually different dimensions the same name.

Args:

  • other (Self): The other dimension to check equality against.

Returns:

True if the dimensions have the same name, False otherwise.

StaticDim

A static tensor dimension.

Static tensor dimensions will always have exactly the same value, and are key to good model performance.

Static dimensions can be created implicitly in most cases: MOTensor(DType.int64, 4, 5) is a tensor with 2 static dimensions, 4 and 5 respectively.

Fields:

  • dim (SIMD[si64, 1]): The size of the static dimension.

Implemented traits:

AnyType, CollectionElement, Copyable, Movable

Methods:

__init__

__init__(inout self: Self, dim: Int)

Int conversion constructor.

Args:

  • dim (Int): The size of the static dimension.

__eq__

__eq__(self: Self, other: Self) -> Bool

Whether the dimension has the same size as another dimension.

Args:

  • other (Self): The other dimension to check equality against.

Returns:

True if both dimensions have the same static size, False otherwise.

Dim

A tensor dimension.

Tensor dimensions can be

  • Static, aka known size
  • Dynamic, aka unknown size
  • Symbolic, aka unknown size but named

In most cases you don't need to work with a Dim directly, but can rely on conversion constructors, for instance you can specify a tensor type as

from max.graph import Dim, MOTensor
var tensor_type = MOTensor(DType.int64, "batch", 10, Dim.dynamic())

will create a tensor type with 3 dimensions: a symbolic "batch" dimension, a static dimension of size 10, and a dynamic dimension.


You can still construct dimensions explicitly via helpers, eg.

```mojo
var some_dims = [
Dim.dynamic(),
Dim.symbolic("batch"),
Dim.static(5),
]

Constraining tensor dimensions is one important way to improve model performance. If tensors have unknown dimensions, we can't optimize them as aggressively. Symoblic tensors allow the compiler to learn constraints on a specific dimension (eg. if 2 inputs have the same batch dimension) which can be an important improvement over dynamic dimensions, but static dims are the easiest to optimize and therefore the easiest to create and work with.

Fields:

  • value (Variant[DynamicDim, StaticDim, SymbolicDim]): The dimension data.

Implemented traits:

AnyType, CollectionElement, Copyable, Movable

Methods:

__init__

__init__(inout self: Self, dim: Int)

Int static dimension conversion constructor.

Args:

  • dim (Int): The static size of the dimension.

__init__(inout self: Self, name: StringLiteral)

StringLiteral symbolic dimension conversion constructor.

Args:

  • name (StringLiteral): The name of the symbolic dimension.

__eq__

__eq__(self: Self, other: Self) -> Bool

Checks whether two dimensions are equal.

Dimensions are equal if they are the same dimension type (dynamic, symbolic, static). Additionally, static dimensions are only equal if their dimension is the same size, and symbolic dimensions are only equal if they have the same name.

Args:

  • other (Self): The other dimension to check equality against.

Returns:

True if the dimensions are equal, False otherwise.

__ne__

__ne__(self: Self, other: Self) -> Bool

Checks whether two dimensions are not equal.

The inverse of eq.

Args:

  • other (Self): The other dimension to check inequality against.

Returns:

False if the dimensions are equal, True otherwise.

static

static static(dim: SIMD[si64, 1]) -> Self

Explicitly constructs a static dimension.

Args:

  • dim (SIMD[si64, 1]): The static size of the dimension.

Returns:

A static dimension of size dim.

symbolic

static symbolic(name: String) -> Self

Explicitly constructs a symbolic dimension.

Args:

  • name (String): The unique name of the dimension.

Returns:

A symbolic dimension with the given name.

dynamic

static dynamic() -> Self

Explicitly constructs a dynamic dimension.

Returns:

A dynamic dimension.

is_dynamic

is_dynamic(self: Self) -> Bool

Checks whether or not the dimension is a dynamic dimension.

Returns:

True if the dimension is dynamic, False otherwise.

is_static

is_static(self: Self) -> Bool

Checks whether or not the dimension is a static dimension.

Returns:

True if the dimension is static, False otherwise.

is_symbolic

is_symbolic(self: Self) -> Bool

Whether or not the dimension is a symbolic dimension.

Returns:

True if the dimension is symbolic, False otherwise.

num_elements

num_elements(self: Self) -> SIMD[si64, 1]

Returns the number of elements in the dimension, if known.

Returns:

For a static dimension, we return the known static dimension size. Otherwise, return an internal value representing an unknown dimension size.

to_mlir

to_mlir(self: Self, m: Module) -> Attribute

Creates an _mlir.Attribute representing this dimension.

This is used internally when constructing tensor _mlir types.

Args:

  • m (Module): A Module instance holding an _mlir.Context.

Returns:

A _mlir.Attribute in the Module's context representing the dimension.

ElementType

The element type of a data container, like a tensor or scalar.

Prefer to use the standard library DType and implicitly convert to this type rather than using it directly.

Fields:

  • dtype (DType): The underlying dtype.

Implemented traits:

AnyType, Copyable, MOType, Movable

Methods:

to_mlir

to_mlir(self: Self, m: Module) -> Type

Converts to an _mlir.Type instance.

Args:

  • m (Module): The Module object holding an _mlir.Context to create in.

Returns:

An _mlir.Type in the specified Context.

to_string

to_string(self: Self, m: Module) -> String

Converts to a maybe-human-readable string.

Args:

  • m (Module): A module object to help with string construction.

Returns:

A string representation of the type.

MOTensor

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 (ElementType): The element type of the tensor value.
  • dims (List[Dim]): The dimensions of the tensor value, if it is known-rank.
  • ranked (Bool): Whether the tensor has a known static rank or not.

Implemented traits:

AnyType, CollectionElement, Copyable, MOType, Movable

Methods:

__init__

__init__(inout self: Self, dtype: ElementType, *dims: Dim)

Constructs a ranked tensor type.

Args:

  • dtype (ElementType): 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: ElementType, ranked: Bool)

Constructs a fully dynamic tensor or 0-dimensional tensor type.

Args:

  • dtype (ElementType): The element type of the tensor data.
  • ranked (Bool): If False, create a fully dynamic tensor. If True, create a rank 0 tensor. This is the same as calling the constructor with just a dtype argument.

__init__(inout self: Self, dtype: ElementType, dim: Int)

Constructs a rank-1 static tensor type.

This is a temporary overload that exists to prevent an overload ambiguity via implicit conversion to a DynamicTensor. It's functionally the same as the ranked tensor variadic constructor.

Args:

  • dtype (ElementType): The element type of the tensor data.
  • dim (Int): The static shape of the singular dimension.

__init__(inout self: Self, dtype: ElementType, dims: List[Dim])

Constructs a ranked tensor type.

Args:

  • dtype (ElementType): 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, m: Module) -> Type

Converts to an _mlir.Type instance.

Args:

  • m (Module): The Module object holding an _mlir.Context to create in.

Returns:

An _mlir.Type in the specified Context.

to_string

to_string(self: Self, m: Module) -> String

Converts to a maybe-human-readable string.

Args:

  • m (Module): A module object to help with string construction.

Returns:

A string representation of the type.

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.

This is not the same as ranked. A tensor must be both ranked and 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.

Raises: If the dimension is out-of-bounds, or if the tensor is unranked.

Args:

  • pos (Int): The dimension index to retrieve.

Returns:

The dimension value at dimension pos.

num_elements

num_elements(self: Self) -> SIMD[si64, 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, MOTensor doesn't actually hold any element values at all.

For any non-static tensor, ie. a tensor having dynamic rank or 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: ElementType) -> Self

Constructs a new tensor type of the same shape with the new dtype.

Args:

  • dtype (ElementType): The new element type for the tensor.

Returns:

A new tensor type with the same shape, and the new element type.

MOList

A type representing a flat list of tensor values.

This isn't an eager list type! It doesn't contain any data, but represents a runtime list that contains tensors.

Fields:

  • eltype (MOTensor): The tensor type of elements in the list.

Implemented traits:

AnyType, CollectionElement, Copyable, MOType, Movable

Methods:

to_mlir

to_mlir(self: Self, m: Module) -> Type

Converts to an _mlir.Type instance.

Args:

  • m (Module): The Module object holding an _mlir.Context to create in.

Returns:

An _mlir.Type in the specified Context.

to_string

to_string(self: Self, m: Module) -> String

Converts to a maybe-human-readable string.

Args:

  • m (Module): A module object to help with string construction.

Returns:

A string representation of the type.

AnyMOType

Represents any possible type for Graph Symbol values.

Every Symbol has a Type, and that type is represented by an AnyMOType. This type may be inspected to get finer-grained types and learn more about an individual Value.

Fields:

  • type (Variant[MOTensor, MOList]): The type data.

Implemented traits:

AnyType, CollectionElement, Copyable, MOType, Movable

Methods:

__init__

__init__(inout self: Self, t: MOTensor)

Constructs a type from a tensor type.

Args:

  • t (MOTensor): The tensor type.

__init__(inout self: Self, t: MOList)

Constructs a type from a list type.

Args:

  • t (MOList): The list type.

list

list(self: Self) -> MOList

Extracts the type as a list type.

This doesn't have any impact at graph execution time, it just retrieves the underlying list type for a type which is a list.

Raises: If the type is some other data type besides a list.

Returns:

The underlying type specifically as a list type.

tensor

tensor(self: Self) -> MOTensor

Extracts the type as a tensor type.

This doesn't have any impact at graph execution time, it just retrieves the underlying tensor type for a type which is a tensor.

Raises: If the type is some other data type besides a tensor.

Returns:

The underlying type specifically as a tensor type.

to_mlir

to_mlir(self: Self, m: Module) -> Type

Converts to an _mlir.Type instance.

Args:

  • m (Module): The Module object holding an _mlir.Context to create in.

Returns:

An _mlir.Type in the specified Context.

from_mlir

static from_mlir(t: Type) -> Self

Constructs a type from an _mlir type.

Args:

  • t (Type): The _mlir Type object to parse into a type.

Returns:

The type represented by the _mlir Type value.

to_string

to_string(self: Self, m: Module) -> String

Converts to a maybe-human-readable string.

Args:

  • m (Module): A module object to help with string construction.

Returns:

A string representation of the type.

TypeTuple

A sequence of 0 or more types.

This is a helper type for graph construction.

Fields:

  • elts (List[AnyMOType]): The sequence of types.

Implemented traits:

AnyType, Copyable, Movable, Sized

Methods:

__init__

__init__(inout self: Self, *elts: AnyMOType)

Constructs a TypeTuple from any number of types.

Args:

  • elts (*AnyMOType): The sequence of types.

__init__(inout self: Self, t: MOTensor)

Constructs a 1-element TypeTuple from a tensor type.

Args:

  • t (MOTensor): The tensor type.

__init__(inout self: Self, t: MOList)

Constructs a 1-element TypeTuple from a list type.

Args:

  • t (MOList): The list type.

__len__

__len__(self: Self) -> Int

Gets the length of the tuple.

Returns:

The number of elements in the tuple.

to_mlir

to_mlir(self: Self, m: Module) -> List[Type]

Converts to a sequence of _mlir.Type instances.

Args:

  • m (Module): The Module object holding an _mlir.Context to create in.

Returns:

A list of _mlir.Types representing the tuple's types.

append

append(inout self: Self, type: AnyMOType)

Appends a type to the back of the tuple.

Args:

  • type (AnyMOType): The type to add to the tuple.

MOType

An internal helper trait for _mlir construction.

MOTypes have methods to help us convert between our structured types and their _mlir representations.

Implemented traits:

AnyType

Methods:

__del__

__del__(owned self: T, /)

Destroy the contained value.

The destructor receives an owned value and is expected to perform any actions needed to end the lifetime of the object. In the simplest case, this is nothing, and the language treats the object as being dead at the end of this function.

to_mlir

to_mlir(self: T, m: Module) -> Type

Converts to an _mlir.Type instance.

Args:

  • m (Module): The Module object holding an _mlir.Context to create in.

Returns:

An _mlir.Type in the specified Context.

to_string

to_string(self: T, m: Module) -> String

Converts to a maybe-human-readable string.

Args:

  • m (Module): A module object to help with string construction.

Returns:

A string representation of the type.