Python module
dim
Library for graph dimension types.
AlgebraicDim
class max.graph.dim.AlgebraicDim(value)
An algebraic tensor dimension to enable expressions over symbolic dimensions.
That is, any expression over a symbolic dimension returns AlgebraicDim
.
Furthermore, algebraic dimensions automatically simplify into a canonical
form.
The following example demonstrates how to create and use algebraic dimensions with symbolic values:
from max.graph import AlgebraicDim, Dim
isinstance(Dim("batch") * 5, AlgebraicDim) # Returns True
print(Dim("batch") * 5) # Outputs: batch * 5
-Dim("x") - 4 == -(Dim("x") + 4) # Returns True
from max.graph import AlgebraicDim, Dim
isinstance(Dim("batch") * 5, AlgebraicDim) # Returns True
print(Dim("batch") * 5) # Outputs: batch * 5
-Dim("x") - 4 == -(Dim("x") + 4) # Returns True
Converts valid input values to Dim.
-
Parameters:
-
attr (
ParamOperatorAttr
)
apply()
classmethod apply(op, *operands)
attr
attr: ParamOperatorAttr
from_mlir()
static from_mlir(attr)
Constructs a dimension from an mlir.Attribute
.
-
Parameters:
-
- dim_attr – The MLIR Attribute object to parse into a dimension.
- attr (
TypedAttr
)
-
Returns:
-
The dimension represented by the MLIR Attr value.
-
Return type:
parameters
property parameters: Iterable[SymbolicDim]
Lists the symbolic dimension names on which this dim depends.
to_mlir()
to_mlir()
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.
-
Return type:
-
ParamOperatorAttr
Dim
class max.graph.dim.Dim(value)
A tensor dimension.
Tensor dimensions can be one of three types:
- Static: Known size
- Symbolic: Unknown size but named
- Algebraic: Unknown size has an algebraic expression
In most cases, you don’t need to work with a Dim
directly.
Instead, use conversion constructors:
from max.graph import Dim, TensorType, DeviceRef
tensor_type = TensorType(DType.int64, ("batch", 10), device=DeviceRef.CPU())
from max.graph import Dim, TensorType, DeviceRef
tensor_type = TensorType(DType.int64, ("batch", 10), device=DeviceRef.CPU())
This creates a tensor type with three dimensions:
- A symbolic “batch” dimension
- A static dimension of size 10
For explicit dimension construction, use the following helpers:
from max.graph import Dim
some_dims = [
SymbolicDim("batch"),
StaticDim(5),
AlgebraicDim(Dim("batch") + 1),
]
from max.graph import Dim
some_dims = [
SymbolicDim("batch"),
StaticDim(5),
AlgebraicDim(Dim("batch") + 1),
]
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), but static dims are the easiest to optimize and therefore the easiest to create and work with.
Converts valid input values to Dim.
-
Parameters:
-
value (
DimLike
)
from_mlir()
static from_mlir(attr)
Constructs a dimension from an mlir.Attribute
.
-
Parameters:
-
- dim_attr – The MLIR Attribute object to parse into a dimension.
- attr (
TypedAttr
)
-
Returns:
-
The dimension represented by the MLIR Attr value.
-
Return type:
parameters
property parameters: Iterable[SymbolicDim]
Lists the symbolic dimension names on which this dim depends.
to_mlir()
to_mlir()
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. -
Return type:
-
TypedAttr
StaticDim
class max.graph.dim.StaticDim(value)
A static tensor dimension.
Static tensor dimensions will always have exactly the same value, and are key to good model performance.
The following example shows how static dimensions can be created implicitly:
from max.graph import TensorType
from max.dtype import DType
tensor = TensorType(DType.int64, (4, 5))
from max.graph import TensorType
from max.dtype import DType
tensor = TensorType(DType.int64, (4, 5))
Converts valid input values to Dim.
-
Parameters:
-
dim (
int
)
dim
dim: int
The size of the static dimension.
from_mlir()
static from_mlir(attr)
Constructs a dimension from an mlir.Attribute
.
-
Parameters:
-
- dim_attr – The MLIR Attribute object to parse into a dimension.
- attr (
TypedAttr
)
-
Returns:
-
The dimension represented by the MLIR Attr value.
-
Return type:
parameters
property parameters: Iterable[SymbolicDim]
Lists the symbolic dimension names on which this dim depends.
to_mlir()
to_mlir()
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. -
Return type:
-
IntegerAttr
SymbolicDim
class max.graph.dim.SymbolicDim(value)
A symbolic tensor dimension.
Symbolic dimensions represent named dimensions in MO tensor types.
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.
In tensor type notation:
!mo.tensor<[batch, x, 10], si32]>
!mo.tensor<[batch, x, 10], si32]>
The first and second dimensions are named batch
and x
respectively.
Creating a SymbolicDim
:
dim = SymbolicDim("name")
dim = SymbolicDim("name")
Using SymbolicDim
in a TensorType
:
tensor_type = TensorType(DType.bool, (SymbolicDim("batch"), SymbolicDim("x"), 10))
tensor_type = TensorType(DType.bool, (SymbolicDim("batch"), SymbolicDim("x"), 10))
Converts valid input values to Dim.
-
Parameters:
-
name (
str
)
from_mlir()
static from_mlir(attr)
Constructs a dimension from an mlir.Attribute
.
-
Parameters:
-
- dim_attr – The MLIR Attribute object to parse into a dimension.
- attr (
TypedAttr
)
-
Returns:
-
The dimension represented by the MLIR Attr value.
-
Return type:
name
name: str
The name of the dimension.
parameters
property parameters: Iterable[SymbolicDim]
Lists the symbolic dimension names on which this dim depends.
to_mlir()
to_mlir()
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. -
Return type:
-
ParamDeclRefAttr
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!