Mojo struct
Dim
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())
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),
]
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
,
UnknownDestructibility
Methods
__init__
__init__(out self, dim: Int)
Int static dimension conversion constructor.
Args:
- dim (
Int
): The static size of the dimension.
__init__(out self, name: StringLiteral)
StringLiteral symbolic dimension conversion constructor.
Args:
- name (
StringLiteral
): The name of the symbolic dimension.
__eq__
__eq__(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, 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) -> 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) -> 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) -> Bool
Whether or not the dimension is a symbolic dimension.
Returns:
True if the dimension is symbolic, False otherwise.
num_elements
num_elements(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.
maybe_num_elements
maybe_num_elements(self) -> Optional[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 None.
to_mlir
to_mlir(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.
from_mlir
static from_mlir(dim_attr: Attribute) -> Self
Constructs a dimension from an _mlir Attribute.
Args:
- dim_attr (
Attribute
): The _mlir Attribute object to parse into a dimension.
Returns:
The dimension represented by the _mlir Attr value.
__str__
__str__(self) -> String
Creates a string representation of the dimension.
Returns:
A human-readable string of the dimension.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!