Skip to main content

struct

Dim

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
var tensor_type = TensorType(DType.int64, "batch", 10, Dim.dynamic())

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.

```mojo
var some_dims = [
Dim.dynamic(),
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. Symoblic 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.

Fields

  • value (Variant[DynamicDim, StaticDim, SymbolicDim]): The dimension data.

Implemented traits

AnyType, CollectionElement, Copyable, Movable, Stringable

Methods

__init__

__init__(inout self: Self, dim: Int)

Int static dimension conversion constructor.

Args:

  • dim (Int): The static size of the dimension.

__init__(inout self: Self, name: StringLiteral)

StringLiteral symbolic dimension conversion constructor.

Args:

  • name (StringLiteral): The name of the symbolic dimension.

__eq__

__eq__(self: Self, other: Self) -> Bool

Checks whether two dimensions are equal.

Dimensions are equal if they are the same dimension type (dynamic, symbolic, static). Additionally, static dimensions are only equal if their dimension is the same size, and symbolic dimensions are only equal if they have the same name.

Args:

  • other (Self): The other dimension to check equality against.

Returns:

True if the dimensions are equal, False otherwise.

__ne__

__ne__(self: Self, other: Self) -> Bool

Checks whether two dimensions are not equal.

The inverse of eq.

Args:

  • other (Self): The other dimension to check inequality against.

Returns:

False if the dimensions are equal, True otherwise.

static

static static(dim: SIMD[int64, 1]) -> Self

Explicitly constructs a static dimension.

Args:

  • dim (SIMD[int64, 1]): The static size of the dimension.

Returns:

A static dimension of size dim.

symbolic

static symbolic(name: String) -> Self

Explicitly constructs a symbolic dimension.

Args:

  • name (String): The unique name of the dimension.

Returns:

A symbolic dimension with the given name.

dynamic

static dynamic() -> Self

Explicitly constructs a dynamic dimension.

Returns:

A dynamic dimension.

is_dynamic

is_dynamic(self: Self) -> Bool

Checks whether or not the dimension is a dynamic dimension.

Returns:

True if the dimension is dynamic, False otherwise.

is_static

is_static(self: Self) -> Bool

Checks whether or not the dimension is a static dimension.

Returns:

True if the dimension is static, False otherwise.

is_symbolic

is_symbolic(self: Self) -> Bool

Whether or not the dimension is a symbolic dimension.

Returns:

True if the dimension is symbolic, False otherwise.

num_elements

num_elements(self: Self) -> SIMD[int64, 1]

Returns the number of elements in the dimension, if known.

Returns:

For a static dimension, we return the known static dimension size. Otherwise, return an internal value representing an unknown dimension size.

to_mlir

to_mlir(self: Self, ctx: Context) -> Attribute

Creates an _mlir.Attribute representing this dimension.

This is used internally when constructing tensor _mlir types.

Args:

  • ctx (Context): The mlir.Context in which to create the attribute.

Returns:

A _mlir.Attribute in the context representing the dimension.