Sharding

Picture an array too big for one chip, cut into tiles and spread across eight of them, one tile per device. JAX names that arrangement with two objects. A Mesh names your devices along axes, so a (4, 2) mesh might call one axis data and the other model. A PartitionSpec maps each array dimension onto a mesh axis, or onto nothing, which means that dimension stays whole on every device. Put a mesh and a spec together and you get a NamedSharding: the mesh says what devices exist and how they're arranged, the spec says which array dimension rides on which axis.

jax.device_put(x, sharding) takes a global array and lays it across the mesh according to that sharding. jax.debug.visualize_array_sharding prints the layout back as a grid of device indices, so you can confirm the tiling matches what you intended instead of trusting the code silently.

Every exercise in this chapter runs on eight devices that don't exist. Set XLA_FLAGS=--xla_force_host_platform_device_count=8 before importing JAX and one laptop CPU pretends to be an eight-device pod. The mesh, the sharding, and the collectives XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → inserts are all real; only the hardware underneath is faked, which is what makes sharding a laptop exercise rather than a cluster requirement.

# run with: XLA_FLAGS=--xla_force_host_platform_device_count=8
import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.ones((8, 16)), NamedSharding(mesh, P("data", "model")))
jax.debug.visualize_array_sharding(x)   # 4x2 tiles, one per device

Nothing about the matmuls, the reductions, or the control flow changes when an array is sharded. You keep writing global-array math exactly as chapters 1 through 9 taught it. What changes sits underneath: XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s SPMD partitioner, called GSPMD, looks at the sharding on your inputs and splits every operation across the mesh to match, inserting whatever collectives the math requires along the way. An all-gather where a full row needs assembling. A reduce-scatter where a sum needs splitting back out. An all-reduce where every device needs the same total. Sharding propagates forward from your inputs through the whole program, so you rarely have to state more than where the leaves start.

Write global math; the compiler places it.

Propagation guesses wrong sometimes: an intermediate value ends up sharded in a way that forces an expensive reshuffle right before an operation that wanted it laid out differently. jax.lax.with_sharding_constraint pins that intermediate's layout explicitly, overriding the guess. This is the second level of control: not choosing the collectives yourself, just correcting where propagation would have chosen badly.

Data parallelism, tensor parallelism, FSDP: under this model none of them are frameworks or separate code paths. They are PartitionSpecs. Data parallelism shards the batch axis and replicates the parameters. Tensor parallelism shards a parameter axis and replicates the batch. FSDP shards the parameters too, gathering them back only when a layer needs the whole thing. Naming the pattern doesn't require different code, only a different spec.

GSPMD chooses the collectives for you, and most of the time that's the right amount of control to give up. jax.shard_map gives it back. Call jax.shard_map(f, mesh=mesh, in_specs=..., out_specs=...) and f runs once per device, on that device's local block: the shapes shrink to what one device actually holds, but the rank stays the same, so the function you write still reads like ordinary array code. Inside f you call the collectives yourself, jax.lax.psum, ppermute, and the rest, addressed by mesh axis name rather than by device index.

One import detail is worth stating plainly rather than leaving to guesswork. On JAX versions before shard_map reached the top-level namespace, it lived at jax.experimental.shard_map.shard_map. Same function, same contract, older address.

Reach for shard_map when the partitioner keeps fighting you toward a layout you didn't want, or when you're writing the communication pattern on purpose rather than hoping the compiler infers it. The kernel path's distributed chapter builds collectives like these from raw remote DMAsA chip pushes a buffer straight into a neighbor’s memory and signals a semaphore, while its compute keeps working. The native distributed operation.taught in /l/ici →, one layer further down; this chapter stops at the level where you name the collective and let JAX lower it.

from functools import partial

import jax
from jax.sharding import PartitionSpec as P

@partial(jax.shard_map, mesh=mesh, in_specs=P("data", None), out_specs=P())
def global_mean(block):              # block: this device's (2, 16) slice
    return jax.lax.pmean(block.mean(), axis_name="data")

Two PartitionSpecs for the same array can both run correctly and still cost very different amounts, because the collectives GSPMD inserts move real bytes over real links. Which spec wins is arithmetic: bytes moved divided by link bandwidth, summed over every collective the choice implies.

This chapter doesn't rebuild that arithmetic. The scaling book carries the full treatment of the trade, and the kernel path's ICIThe inter-chip links (4.5e10 bytes per second each way per link on v5e); every collective resolves to hops over these.taught in /l/ici → chapter teaches the mechanism underneath the collectives themselves, down to the remote DMA. What this chapter leaves you with is the vocabulary to read that treatment with: mesh, spec, propagation, and the three levels of control between them.