Python module

type

Library for graph value types.

`AlgebraicDim`

class max.graph.type.AlgebraicDim(value: int | str | Dim | integer)

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.

For example:

>>> from max.graph import AlgebraicDim, Dim
>>> isinstance(Dim("batch") * 5, AlgebraicDim)
True
>>> print(Dim("batch") * 5)
batch * 5
>>> -Dim("x") - 4 == -(Dim("x") + 4)
True
>>> from max.graph import AlgebraicDim, Dim
>>> isinstance(Dim("batch") * 5, AlgebraicDim)
True
>>> print(Dim("batch") * 5)
batch * 5
>>> -Dim("x") - 4 == -(Dim("x") + 4)
True

`attr`

attr*: Attribute*

`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.
Return type:

Dim

`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.

`to_str()`

to_str(use_repr: bool)

`Dim`

class max.graph.type.Dim(value: int | str | Dim | integer)

A tensor dimension.

Tensor dimensions can be one of three types:

Static: Known size
Dynamic: Unknown size
Symbolic: Unknown size but named

In most cases, you don’t need to work with a Dim directly. Instead, use conversion constructors:

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))

This creates a tensor type with three dimensions:

A symbolic “batch” dimension
A static dimension of size 10
A dynamic dimension

For explicit dimension construction, use the following helpers:

from max.graph import Dim

some_dims = [
    Dim.symbolic("batch"),
    Dim.static(5),
]
from max.graph import Dim

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.
Return type:

Dim

`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

`static_dims`

property static_dims*: list[int]*

Returns all static dims in the shape as a list of integers.

`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.

`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.

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"), Dim.dynamic(), 10))
tensor_type = TensorType(DType.bool, (SymbolicDim("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.
Return type:

Dim

`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], device: Device | None = None)

A symbolic TensorType.

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.

An optional device (device) can also be provided to indicate the explicit device the tensor is associated with.

`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, device, and the new element type.

`device`

device*: Device | None*

The device of the tensor value.

`dim()`

dim(pos: int) → 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:

RuntimeError – 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.

`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.

AlgebraicDim​

attr​

from_mlir()​

to_mlir()​

to_str()​

Dim​

from_mlir()​

to_mlir()​

Shape​

rank​

static_dims​

to_mlir()​

StaticDim​

dim​

from_mlir()​

to_mlir()​

SymbolicDim​

from_mlir()​

name​

to_mlir()​

TensorType​

cast()​

device​

dim()​

dtype​

from_mlir()​

num_elements()​

rank​

shape​

to_mlir()​

Type​

from_mlir()​

to_mlir()​

`AlgebraicDim`

`attr`

`from_mlir()`

`to_mlir()`

`to_str()`

`Dim`

`from_mlir()`

`to_mlir()`

`Shape`

`rank`

`static_dims`

`to_mlir()`

`StaticDim`

`dim`

`from_mlir()`

`to_mlir()`

`SymbolicDim`

`from_mlir()`

`name`

`to_mlir()`

`TensorType`

`cast()`

`device`

`dim()`

`dtype`

`from_mlir()`

`num_elements()`

`rank`

`shape`

`to_mlir()`

`Type`

`from_mlir()`

`to_mlir()`