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:

Static, aka known size

Dynamic, aka unknown size

Symbolic, aka unknown size but named

In most cases you don’t need to work with a Dim directly, but can rely on conversion constructors, for instance you can specify a tensor type as

from max.graph import Dim, TensorType
tensor_type = TensorType(DType.int64, ("batch", 10))

will create a tensor type with 3 dimensions: a symbolic “batch” dimension, a static dimension of size 10, and a dynamic dimension.

You can still construct dimensions explicitly via helpers, eg.

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.

`is_static()`

is_static() → bool

Checks whether or not the dimension is a static dimension.

Returns:

True if the dimension is static, False otherwise.

`is_symbolic()`

is_symbolic() → bool

Whether or not the dimension is a symbolic dimension.

Returns:

True if the dimension is symbolic, False otherwise.

`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

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

`is_static()`

is_static() → bool

Checks whether or not the dimension is a static dimension.

Returns:

True if the dimension is static, False otherwise.

`is_symbolic()`

is_symbolic() → bool

Whether or not the dimension is a symbolic dimension.

Returns:

True if the dimension is symbolic, False otherwise.

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

SymbolicDims`s have a name and are printed as their name on MO types, eg. `!mo.tensor<[batch, x, 10], si32]> the first and second dimensions are named “batch” and “x” respectively.

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.

Create a symbolic dimension via SymbolicDim(“name”), for example: TensorType(DType.bool, (“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.

`is_static()`

is_static() → bool

Checks whether or not the dimension is a static dimension.

Returns:

True if the dimension is static, False otherwise.

`is_symbolic()`

is_symbolic() → bool

Whether or not the dimension is a symbolic dimension.

Returns:

True if the dimension is symbolic, False otherwise.

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

A symbolic tensor type.

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.

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

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

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.

`is_static()`

is_static() → bool

Checks whether the tensor type has a fully static shape or not.

A tensor must have all of its dimensions be static (or be 0-dimensional) in order to be static.

Returns:

True if the tensor has a fully static shape, False otherwise.

`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!
If you'd like to share more information, please report an issue on GitHub

😔 What went wrong?

Dim​

from_mlir()​

is_static()​

is_symbolic()​

to_mlir()​

Shape​

rank​

to_mlir()​

StaticDim​

dim​

from_mlir()​

is_static()​

is_symbolic()​

to_mlir()​

SymbolicDim​

from_mlir()​

is_static()​

is_symbolic()​

name​

to_mlir()​

TensorType​

cast()​

dim()​

dtype​

from_mlir()​

is_static()​

num_elements()​

rank​

shape​

to_mlir()​

Type​

from_mlir()​

to_mlir()​

`Dim`

`from_mlir()`

`is_static()`

`is_symbolic()`

`to_mlir()`

`Shape`

`rank`

`to_mlir()`

`StaticDim`

`dim`

`from_mlir()`

`is_static()`

`is_symbolic()`

`to_mlir()`

`SymbolicDim`

`from_mlir()`

`is_static()`

`is_symbolic()`

`name`

`to_mlir()`

`TensorType`

`cast()`

`dim()`

`dtype`

`from_mlir()`

`is_static()`

`num_elements()`

`rank`

`shape`

`to_mlir()`

`Type`

`from_mlir()`

`to_mlir()`