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)

Parameters:

• op (POC)
• operands (int | str | Dim | integer)

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:

Dim

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:

Dim

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:

Dim

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:

Dim

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!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit