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 module

max.experimental.sharding

Defines how a tensor is laid out across a device mesh and how ops reshard it.

Describes, for every op, what redistribution to perform before the op runs. The pipeline is deliberately local: per-op rules over a placement vocabulary (Replicated, Sharded, Partial, Unknown), scored by a single cost model, with the cheapest plan picked at each dispatch. There is no whole-graph trace.

An op redistributes its inputs whenever the plan it picks needs it, so a model can contain collectives it never wrote. Inside an auto_reshard() block, mode="warn" reports each one and "raise" refuses it, both naming the collective and the transfer_to that would replace it:

from max.driver import CPU
from max.experimental.functional import full, matmul, relu
from max.experimental.sharding import (
    DeviceMesh,
    DeviceMapping,
    Sharded,
    auto_reshard,
)

# A simulated two-device mesh (both slots are the same CPU).
mesh = DeviceMesh(
    devices=(CPU(), CPU()), mesh_shape=(2,), axis_names=("tp",)
)

# ``a`` is column-sharded, ``b`` is row-sharded: a @ b contracts the
# sharded dimension, so the product is a partial sum on every device.
a = full([4, 8], 1.0, device=DeviceMapping(mesh, (Sharded(1),)))
b = full([8, 2], 1.0, device=DeviceMapping(mesh, (Sharded(0),)))

# ``relu`` needs the full sum, so its input is allreduced first. This
# block reports that, instead of letting it pass unremarked.
with auto_reshard(mode="warn"):
    y = relu(matmul(a, b))

This module avoids the overloaded word “rank”. A device is one accelerator; a mesh axis is one named dimension of the DeviceMesh grid; a shard is one device’s piece of a tensor; a tensor axis is a dimension of the tensor itself.

Device mesh

DeviceMesh	An N-dimensional logical grid of devices.

Placements

Partial	Every device holds a part of the value, and the value is their sum.
Placement	Abstract base for all placement types.
Replicated	Every device on this mesh axis holds the same copy of the data.
Sharded	Every device on this mesh axis holds a slice along axis.
Unknown	Per-device values on this mesh axis with no known relation.

Layouts

BufferLayout	A buffer a compiled callable may store through.
TensorLayout	A tensor's dtype and global shape, distributed across a device mesh.

Tensor-to-mesh mappings

DeviceMapping	How a tensor is distributed across a device mesh.
NamedMapping	Builds a DeviceMapping from a JAX-style named spec.

Sharding rules

A sharding rule returns the AxisAssignment rows its op accepts; AxisAssignment describes what a rule receives and returns. These functions build common rules and rows.

AxisAssignment	One placement choice that a sharding rule offers for its op.

match_operand_placement	Returns a sharding rule whose results take operand name's placement.
replicated_rows	Returns the one row of an op that runs only on whole tensors.

Resharding

auto_reshard	Sets how ops in this block may change the placement of their inputs.

ALL_TRANSITIONS	Build an immutable unordered collection of unique elements.
DEFAULT_TRANSITIONS	Build an immutable unordered collection of unique elements.
Transition	Built-in immutable sequence.

Exceptions

ConversionError	Raised when a sharding spec conversion would lose information.
ShardingError	Raised when a sharding constraint cannot be satisfied.