How JAX shards a computation, and why that changed

I have eight devices. I want a matmul whose contracting dimension is split across some of them. So I write the matmul normally, say where the two operands live, and compile:

mesh = Mesh(np.array(jax.devices()).reshape(4, 2), ('data', 'model'))
xs = jax.device_put(x, NamedSharding(mesh, P('data', 'model')))   # K split on 'model'
Ws = jax.device_put(W, NamedSharding(mesh, P('model', None)))     # K split on 'model'
compiled = jax.jit(lambda a, b: jnp.tanh(a @ b)).lower(xs, Ws).compile()

I never wrote a collective. The compiled program has one anyway:

%all-reduce = f32[64,1024]{1,0} all-reduce(%ynn_fusion), channel_id=1,
    replica_groups={{0,1},{2,3},{4,5},{6,7}}, use_global_device_ids=true,
    to_apply=%add.clone, frontend_attributes={is_spmd_generated="true"}

Because I sharded the contracting dimension, each device holds a partial product, and the result is only correct after summing across the shard. The compiler worked that out, grouped the eight devices into the four model-pairs, inserted the all-reduce, and tagged it is_spmd_generated="true" so you can tell its work from yours. That is the whole value proposition: you say where the data is, and the compiler derives what communication is required.

The vocabulary, and the one part that changed

Three types carry the model:

• A Mesh is a named grid of devices. Eight devices arranged 4x2, with the axes named data and model. From there on you never name a physical device; you name mesh axes. That indirection is the good idea, because it decouples the logical parallelism from the count of chips you happen to have.
• A PartitionSpec, written P(...), maps each dimension of an array to a mesh axis, or to None for replicated. P('data', 'model') says split dimension 0 across data, dimension 1 across model.
• A NamedSharding is a Mesh plus a PartitionSpec: concretely how this array is laid out on this grid.

The part that changed is a fourth thing, and it is easy to miss because it lives in a keyword argument. Every mesh axis has an axis type: Auto, Explicit, or Manual. The axis type decides whether placement shows up in the type of a traced value, and the two mesh constructors do not agree on the default:

Mesh(np.array(jax.devices()).reshape(4, 2), ('data', 'model')).axis_types
# (Auto, Auto)

jax.make_mesh((4, 2), ('data', 'model')).axis_types
# (Explicit, Explicit)

Mesh is the old low-level constructor and it gives you Auto. jax.make_mesh is the one the docs now steer you to, and it gives you Explicit. Nearly every sharding tutorial written before about a year ago uses the first one.

Auto: GSPMD infers, and the type stays quiet

Start with Auto, because it is the classical behaviour and the baseline for everything else. Shard the inputs and look at what the tracer knows:

mesh = Mesh(np.array(jax.devices()).reshape(4, 2), ('data', 'model'))
xs = jax.device_put(jnp.ones((256, 512)), NamedSharding(mesh, P('data', 'model')))
print(jax.typeof(xs))
# float32[256,512]

Dtype and shape. The array is physically distributed over eight devices right now, and its type says nothing about that. Placement is real but invisible to the front end.

The system that makes it work anyway is GSPMD, and it runs inside XLA as part of jit. It does two jobs. Propagation: given shardings on some values, infer a consistent sharding for every other value, the way a type inferencer propagates types. Partitioning: rewrite the single logical program into the per-device SPMD program, inserting collectives so the sharded computation equals the original.

Explicit: placement in the type

With Explicit axis types, placement stops being an invisible annotation and becomes part of the value's type. The tracer knows where each value lives, and sharding mistakes get caught at trace time rather than by a runtime propagator. This is the change that makes the old complaint obsolete: placement is no longer an annotation hanging off a value, it is part of its type.

For inference engineering, the payoff is a mental model of multi-device serving that matches how the hardware actually works: the mesh is your cluster topology, the PartitionSpec is your tensor layout, and the compiler's collectives are your data movement. When you shard a KV cache or a weight matrix across GPUs, the same machinery applies.

You say where the data is, and the compiler derives what communication is required. That is the whole value proposition.

Back to the blog