Python module
type
Library for graph value types.
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, TensorType
tensor_type = TensorType(DType.int64, ("batch", 10))
from max.graph import Dim, TensorType
tensor_type = TensorType(DType.int64, ("batch", 10))
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.
some_dims = [
Dim.symbolic("batch"),
Dim.static(5),
]
some_dims = [
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. Symbolic 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.
from_mlir()
static from_mlir(dim_attr: Attribute) → Dim
Constructs a dimension from an mlir.Attribute.
-
Parameters:
dim_attr – The MLIR Attribute object to parse into a dimension.
-
Returns:
The dimension represented by the MLIR Attr value.
is_static()
is_static() → bool
Checks whether or not the dimension is a static dimension.
-
Returns:
True if the dimension is static, False otherwise.
is_symbolic()
is_symbolic() → bool
Whether or not the dimension is a symbolic dimension.
-
Returns:
True if the dimension is symbolic, False otherwise.
to_mlir()
to_mlir() → Attribute
Creates an mlir.Attribute representing this dimension.
This is used internally when constructing tensor MLIR types.
-
Returns:
An mlir.Attribute in the context representing the dimension.
Shape
class max.graph.type.Shape(dims: Iterable[int | str | Dim | integer] = ())
rank
property rank
to_mlir()
to_mlir() → Attribute
StaticDim
class max.graph.type.StaticDim(value: int | str | Dim | integer)
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: TensorType(DType.int64, (4, 5)) is a tensor with 2 static dimensions: 4 and 5 respectively.
dim
dim*: int*
The size of the static dimension.
from_mlir()
static from_mlir(dim_attr: Attribute) → Dim
Constructs a dimension from an mlir.Attribute.
-
Parameters:
dim_attr – The MLIR Attribute object to parse into a dimension.
-
Returns:
The dimension represented by the MLIR Attr value.
is_static()
is_static() → bool
Checks whether or not the dimension is a static dimension.
-
Returns:
True if the dimension is static, False otherwise.
is_symbolic()
is_symbolic() → bool
Whether or not the dimension is a symbolic dimension.
-
Returns:
True if the dimension is symbolic, False otherwise.
to_mlir()
to_mlir() → Attribute
Creates an mlir.Attribute representing this dimension.
This is used internally when constructing tensor MLIR types.
-
Returns:
An mlir.Attribute in the context representing the dimension.
SymbolicDim
class max.graph.type.SymbolicDim(value: int | str | Dim | integer)
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 SymbolicDim(“name”), for example: TensorType(DType.bool, (“batch”, Dim.dynamic(), 10)).
from_mlir()
static from_mlir(dim_attr: Attribute) → Dim
Constructs a dimension from an mlir.Attribute.
-
Parameters:
dim_attr – The MLIR Attribute object to parse into a dimension.
-
Returns:
The dimension represented by the MLIR Attr value.
is_static()
is_static() → bool
Checks whether or not the dimension is a static dimension.
-
Returns:
True if the dimension is static, False otherwise.
is_symbolic()
is_symbolic() → bool
Whether or not the dimension is a symbolic dimension.
-
Returns:
True if the dimension is symbolic, False otherwise.
name
name*: str*
The name of the dimension.
to_mlir()
to_mlir() → Attribute
Creates an mlir.Attribute representing this dimension.
This is used internally when constructing tensor MLIR types.
-
Returns:
An mlir.Attribute in the context representing the dimension.
TensorType
class max.graph.type.TensorType(dtype: DType, shape: Iterable[int | str | Dim | integer])
A symbolic tensor type.
This is _not_ an eager tensor type! This contains no actual data, but instead represents the type of 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 optimization 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.
cast()
cast(dtype: DType) → TensorType
Constructs a new tensor type of the same shape with the new dtype.
-
Parameters:
dtype – The new element type for the tensor.
-
Returns:
A new tensor type with the same shape, and the new element type.
dim()
Gets the pos’th dimension of the tensor type.
Supports negative-indexing, ie. t.dim(-1) will give the last dimension.
-
Parameters:
pos – The dimension index to retrieve.
-
Returns:
The dimension value at dimension pos.
-
Raises:
If the dimension is out-of-bounds. –
dtype
dtype*: DType*
The element type of the tensor value.
from_mlir()
static from_mlir(t: Type) → TensorType
Constructs a tensor type from an MLIR type.
-
Parameters:
t – The MLIR Type object to parse into a tensor type.
-
Returns:
The tensor type represented by the MLIR Type value.
is_static()
is_static() → 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.
num_elements()
num_elements() → int
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, in other words a tensor having any symbolic dimensions, the return value will be meaningless.
-
Returns:
The number of elements the tensor contains.
rank
property rank*: int*
Gets the rank of the tensor type.
-
Returns:
The tensor’s static rank.
shape
shape*: Shape*
The dimensions of the tensor value.
to_mlir()
to_mlir() → Type
Converts to an mlir.Type instance.
-
Returns:
An mlir.Type in the specified Context.
Type
class max.graph.type.Type
Represents any possible type for Graph values.
Every Value in the Graph has a Type, and that type is represented by an Type. This type may be inspected to get finer-grained types and learn more about an individual Value.
from_mlir()
static from_mlir(t: Type) → Type
Constructs a type from an MLIR type.
-
Parameters:
t – The MLIR Type object to parse into a type.
-
Returns:
The type represented by the MLIR Type value.
to_mlir()
to_mlir() → Type
Converts to an mlir.Type instance.
-
Returns:
An mlir.Type in the specified Context.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!