For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

PerShardDim

PerShardDim

class max.tensor.sharding.per_shard_dim.PerShardDim(per_shard=(), *, global_dim=None)

source

Bases: Dim

A Dim listing one cell per mesh shard.

Appears on a Sharded axis when the cells differ, and also when a rule keeps the wrapper to record the axis’s global extent (force_wrap=True or an attached global_dim), which is what shape does for every sharded axis. Test it with is_per_shard_dim() rather than by comparing cells. The wrapper must be projected per shard via local_dim_at() before reaching MLIR.

Parameters:

• per_shard (tuple[Dim, ...]) – One cell per device, in mesh order.
• global_dim (Dim | None) – The global extent this wrapper records, when it has one.

Return type:

PerShardDim

Allocates the wrapper, returning per_shard itself on a plain re-wrap.

is_static

property is_static: bool

source

True if this axis’s global extent is a static size.

Folds to the global dim first (see global_dim()), so a sharded axis whose global is static reports True even though isinstance(self, StaticDim) is False.

is_symbolic

property is_symbolic: bool

source

True if this axis’s global extent is a symbolic (named) dim.

parameters

property parameters: Iterable[SymbolicDim]

source

Distinct symbolic-dim parameters referenced across all cells.

per_shard

per_shard: tuple[Dim, ...]

source

to_mlir()

to_mlir()

source

Raises; wrappers must be projected per shard before reaching MLIR.

Return type:

NoReturn