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 allowed row picked on each mesh axis. A dispatch whose every row needs a transition the current policy disallows raises rather than picking one. 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`` has no row for a partial sum, so its input moves to Replicated.
# The default transition set makes that an allreduce: the cheaper
# reduce-scatter into Sharded is not allowed outside an
# ``auto_reshard(ALL_TRANSITIONS)`` block. This block reports the move.
with auto_reshard(mode="warn"):
    y = relu(matmul(a, b))

This module avoids the overloaded word “rank” in prose. 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. TensorLayout.rank is the number of tensor axes.

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 receives the op’s arguments with each tensor replaced by its TensorLayout, and returns the AxisAssignment rows the op accepts: what each tensor operand must hold and what the op produces. 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	Every transition the picker can plan with.
DEFAULT_TRANSITIONS	Every transition except (Partial, Sharded), so a partial sum resolves to Replicated by allreduce; ALL_TRANSITIONS also allows reduce-scatter.
Transition	A placement change on one mesh axis, as (from, to) placement types.

Visualization

visualize	Draws a mesh or a tensor's slices as blocks beside the tree they share.

Exceptions

ConversionError	Raised when a named spec cannot be represented on its mesh.
ShardingError	Raised when a sharding rule cannot place a call.