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

Sharded

Sharded

class max.experimental.sharding.Sharded(axis)

source

Bases: Placement

Every device on this mesh axis holds a slice along axis.

Parameters:

axis (int) – The tensor axis along which data is split.

axis

axis: int

source

local_dim()

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

source

Splits parent along mesh_axis into per-shard cells.

parent is split once into mesh.axis_size(mesh_axis) balanced pieces, and the result holds one cell per device of the mesh, in row-major device order. A mesh axis that does not span the whole mesh gives the same piece to every device that shares its coordinate along the splitting axes, and make_per_shard_dim collapses equal cells to a plain dim. For example, 10 rows split over both axes of a 2x2 mesh give devices 0 to 3 sizes 3, 3, 2, and 2.

StaticDim parents use uneven divmod; SymbolicDim parents mint fresh per-shard cells; AlgebraicDim parents raise. Wrapper parents pass through verbatim. When allow_symbolic_mint is False, a bare SymbolicDim raises instead of minting.

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 instead of minting per-shard cells. Defaults to True.

Raises:

ShardingError – If parent is an AlgebraicDim, or a bare SymbolicDim while allow_symbolic_mint is False.

Return type:

Dim

localized_axis()

localized_axis()

source

Returns the tensor axis this Sharded localizes.

Return type:

int | None

transition_to()

transition_to(other)

source

Returns the Collective for the move to another placement.

Replicated takes an allgather. A different Sharded axis returns Collective.ALL_TO_ALL, a cost and not a kernel: transfer_to() allgathers and then splits locally, so the whole tensor crosses the mesh. Partial has no redistribution and reports Collective.INFEASIBLE.

Parameters:

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

Returns:

The Collective that performs the move.

Return type:

Collective