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

Python module

max.experimental.sharding

Distributed-tensor sharding: how a tensor is laid out across a device mesh.

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), scored by a single cost model, with one pluggable Solver making the choice at each dispatch. There is no whole-graph trace.

A mode(...) block selects the solver for the ops inside it:

import numpy as np
from max.driver import CPU
from max.dtype import DType
from max.experimental.functional import full, matmul, relu, transfer_to
from max.experimental.sharding import (
    DeviceMesh,
    GreedyReshard,
    PlacementMapping,
    Sharded,
    mode,
)

# 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 picker resolves the result back to replicated.
a = transfer_to(
    full([4, 8], 1.0, dtype=DType.float32, device=mesh.devices[0]),
    PlacementMapping(mesh, (Sharded(1),)),
)
b = transfer_to(
    full([8, 2], 1.0, dtype=DType.float32, device=mesh.devices[0]),
    PlacementMapping(mesh, (Sharded(0),)),
)

with mode(GreedyReshard(on_reshard="warn")):
    y = relu(matmul(a, b))

Shipped solvers: GreedyReshard (cheapest feasible action), NoReshard (passthrough only; errors on any reshard), and PartialsOnly (only Partial -> Replicated resolutions).

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.

`get_active_mesh`	Returns the mesh from the current `mesh_context()`, or `None`.
`mesh_context`	Publishes `mesh` to spec-first `NamedMapping` constructions.

Placements

`Partial`	Every device holds a partial result that must be reduced.
`Placement`	Abstract base for all placement types.
`ReduceOp`	Reduction operations for partial placements.
`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`.
`Collective`	The collectives the cost model understands.

Distributed types

`DistributedTensorType`	A symbolic type for a tensor distributed across a device mesh.
`TensorLayout`	Metadata snapshot of a distributed tensor for rule evaluation.

Per-op decisions

`Action`	A rule's picked decision for one op call.
`ActionSet`	A rule's menu of per-axis sharding options for one op call.
`AxisAssignment`	One per-mesh-axis row in an `ActionSet`.
`PerShard`	A distinct value per mesh shard.

Pickers

`GreedyReshard`	Default per-op picker: enumerate → cheapest.
`NoReshard`	Passthrough-only picker: returns the first zero-reshard action.
`PartialsOnly`	Cost-model-free picker that only resolves `Partial → Replicated`.
`ReshardBehavior`	alias of `Literal`['silent', 'warn', 'raise']
`Solver`	alias of `Callable`[[`ActionSet`, `Sequence`[`TensorLayout`]], `Action`]

Exceptions

`ShardingError`	Raised when a sharding constraint cannot be satisfied.

Functions

`build_action_set`	Wraps rule-emitted rows into a feasibility-filtered `ActionSet`.
`force_replicated_action_set`	Single-row `(R,…,R) -> R` `ActionSet` for ops that do not expose sharding.
`mode`	Binds `solver` for the duration of a `with` block or function call.

Device mesh​

Placements​

Distributed types​

Per-op decisions​

Pickers​

Exceptions​

Functions​