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

Placement

Placement

class max.experimental.sharding.Placement

source

Bases: ABC

Abstract base for all placement types.

Each placement describes what a single mesh axis does to a tensor: Replicated (full copy), Sharded (split along a tensor dim), Partial (partial result needing reduction), or Unknown (per-device values with no known relation).

local_dim()

local_dim(parent, mesh, mesh_axis, *, allow_symbolic_mint=True)

source

Returns the per-shard local cells of parent along mesh axes.

The default returns parent unchanged; Sharded overrides to split. A wrapper parent is passed through verbatim.

Parameters:

• parent (int | str | Dim | integer | TypedAttr) – The global dim to localize.
• mesh (DeviceMesh) – The device mesh.
• mesh_axis (int | Sequence[int]) – The mesh axis that splits parent, or every mesh axis that splits it.
• allow_symbolic_mint (bool) – When False, a bare SymbolicDim raises ShardingError instead of minting one symbol per shard. MAX callers use the default.

Return type:

Dim

localized_axis()

localized_axis()

source

Tensor axis this placement localizes, or None if it doesn’t.

Return type:

int | None

transition_to()

transition_to(other)

source

Returns the Collective for the move self -> other.

The base returns Collective.NOOP when the two placements are equal and Collective.INFEASIBLE otherwise, so a placement subclass that does not override this can never be moved by the picker.

Parameters:

other (Placement) – The placement this one moves to.

Returns:

The Collective that performs the move.

Return type:

Collective