type
Library for graph Symbol Types.
DynamicDim
A dynamic tensor dimension.
DynamicDim
s 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.
SymbolicDims
s 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.