Skip to main content
Product
Developers
Connect
Company
Log inDownload

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.