Python module

type

Library for graph value types.

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.

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