Python module

type

Library for graph value types.

AlgebraicDim

class max.graph.type.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

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

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

Shape

class max.graph.type.Shape(dims=())

Parameters:

dims (ShapeLike )

from_mlir()

classmethod from_mlir(attr)

Parameters:

attr (TypedAttr )

Return type:

Shape

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

Return type:

ShapeAttr

StaticDim

class max.graph.type.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

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

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

TensorType

class max.graph.type.TensorType(dtype, shape, device)

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.

The following example shows how to create a tensor type with static dimensions and access its properties:

from max.graph import TensorType
from max.dtype import DType
# Create a tensor type with float32 elements and static dimensions 2x3
tensor_type = TensorType(DType.float32, (2, 3))
print(tensor_type.dtype)  # Outputs: DType.float32
print(tensor_type.shape)  # Outputs: [2, 3]

from max.graph import TensorType
from max.dtype import DType
# Create a tensor type with float32 elements and static dimensions 2x3
tensor_type = TensorType(DType.float32, (2, 3))
print(tensor_type.dtype)  # Outputs: DType.float32
print(tensor_type.shape)  # Outputs: [2, 3]

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.

Constructs a tensor type.

Parameters:

• dtype (DType ) – The element type of the tensor data.
• dims – The shape dimensions of the tensor. The number of dims is the rank of the tensor.
• shape (Shape )
• device (DeviceRef )

as_buffer()

as_buffer()

Returns the analogous buffer type.

Return type:

BufferType

from_mlir()

classmethod from_mlir(type)

Constructs a tensor type from an MLIR type.

Parameters:

• t – The MLIR Type object to parse into a tensor type.
• type (TensorType )

Returns:

The tensor type represented by the MLIR Type value.

Return type:

TensorType

to_mlir()

to_mlir()

Converts to an mlir.Type instance.

Returns:

An mlir.Type in the specified Context.

Return type:

TensorType

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.

The following example shows how to work with types in a graph:

from max.graph import Graph, TensorType
from max.dtype import DType
with Graph() as g:
    # Create a tensor constant with a specific type
    tensor_type = TensorType(DType.float32, [2, 3])
    # The type can be inspected to get information about the value
    print(f"Tensor element type: {tensor_type.dtype}")  # Outputs: DType.float32
    print(f"Tensor shape: {tensor_type.shape}")  # Outputs: [2, 3]

from max.graph import Graph, TensorType
from max.dtype import DType
with Graph() as g:
    # Create a tensor constant with a specific type
    tensor_type = TensorType(DType.float32, [2, 3])
    # The type can be inspected to get information about the value
    print(f"Tensor element type: {tensor_type.dtype}")  # Outputs: DType.float32
    print(f"Tensor shape: {tensor_type.shape}")  # Outputs: [2, 3]

from_mlir()

static from_mlir(t)

Constructs a type from an MLIR type.

Parameters:

t (MlirType ) – The MLIR Type object to parse into a type.

Returns:

The type represented by the MLIR Type value.

Return type:

Type

to_mlir()

to_mlir()

Converts to an mlir.Type instance.

Returns:

An mlir.Type in the specified Context.

Return type:

MlirType

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