What is a Mesh in JAX?

A Mesh is how JAX names the accelerators it will spread work across. Its docstring is blunt about the job: a Mesh "declares the hardware resources available in the scope of this manager." Concretely it is "a multidimensional NumPy array of JAX devices, where each axis of the mesh has a name, e.g. 'x' or 'y'." So you take a flat list of eight devices, reshape it into a 2×4 grid, and name the two axes — `Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))`. Nothing is sharded yet; the mesh is only the coordinate system. Everything else refers to devices by these axis names rather than by hardware id, which is what lets the same program run on a different-shaped cluster by rebuilding just the mesh.

What is a PartitionSpec in JAX?

A PartitionSpec says which array dimension is split across which mesh axis. Its docstring: it is a "Tuple describing how to partition an array across a mesh of devices," where "each element is either None, a string, or a tuple of strings." Position matters — the i-th element of the spec describes the i-th dimension of the array. `PartitionSpec('x', 'y')` "says that the first dimension of data is sharded across x axis of the mesh, and the second dimension is sharded across y axis of the mesh." A `None` leaves that dimension replicated (every device along that axis keeps a full copy), and a tuple like `('x', 'y')` shards one array dimension across two mesh axes at once.

What is NamedSharding in JAX?

NamedSharding is the pairing that turns a mesh and a spec into an actual placement. As the docstring puts it, a NamedSharding "expresses sharding using named axes" and "is a pair of a Mesh of devices and PartitionSpec which describes how to shard an array across that mesh." You build one as `NamedSharding(mesh, P('x', 'y'))` and hand it to `jax.device_put` to lay a real array onto the devices, or attach it to a jit input or output. Because it is expressed in mesh-axis names rather than device ids, the same NamedSharding is meaningful on any mesh with the same axis names.

How does jit decide the sharding of intermediate values?

By propagation. You annotate the shardings you care about — some inputs, sometimes an intermediate pinned with `jax.lax.with_sharding_constraint` — and inside jit the partitioner (GSPMD, and now Shardy) propagates those constraints through the rest of the program, filling in every unspecified value with a sharding it infers from the ops. When a computation needs data that lives on other devices, the same pass inserts the communication (an all-reduce over a summed-away axis, an all-gather to un-shard, a reduce-scatter). The catch is that propagation is heuristic: if "there aren't enough sharding constraints in the program, so a small change in propagation order can affect the final result," you can get a surprising layout or an out-of-memory. The fix is to pin the shardings that matter rather than hope the compiler guesses them.

What is the difference between shard_map and jit auto-sharding?

They sit at opposite ends of who writes the communication. Under jit auto-sharding you program with global arrays, annotate what you can, and let the compiler choose the rest and drop in the collectives — easy and composable, but the compiler's choices can surprise you. shard_map hands control back: it "maps a function over shards of data using a mesh of devices," running your function on each device against its local shard, and "the names of the Mesh can be used in collective communication operations in f" — so you call the all-reduce yourself. You specify in_specs and out_specs the same way, but there is no propagation and no compiler-inserted communication; the per-device shapes and the collectives are explicit. Auto for productivity, shard_map for predictable performance and hand-tuned kernels.

What is Shardy and when does it replace GSPMD in JAX?

Shardy is JAX's next-generation partitioner, "a new partitioning system co-developed by GDM Model Scaling (author of PartIR) and XLA/CoreML teams (author of GSPMD)," which "aims to provide better usability and control to users, and will gradually replace GSPMD and PartIR." Per the migration doc, "after the migration is complete in March 2026, Shardy will be the only partitioner in JAX." You can tell it is on by the log line "Using Shardy for XLA SPMD propagation," and until the cutover it can be turned off with `jax.config.update('jax_use_shardy_partitioner', False)`. Because Shardy uses a different propagation order and conflict-resolution heuristics, it "can sometimes output slightly different results" than GSPMD — which is usually a sign the program was under-constrained rather than that either partitioner is wrong.

Sharding in JAX

A 70-billion-parameter model does not fit on one accelerator, so JAX spreads each array across many — without you rewriting the math. Name your devices into a mesh, say which array dimension rides which mesh axis with a PartitionSpec, and the compiler places the bytes and inserts the communication. We build it up from one device to a full mesh, watch what jit does with a sharding — propagation, the gaps the compiler fills, the collectives it drops in — weigh shard_map against auto-sharding honestly, and face the Shardy cutover. Drawn, computed, and animated.

System design · AI / ML. The source ↗

A free, interactive, animated visual explainer of Sharding in JAX — built to be understood, not skimmed.

Questions

What is a Mesh in JAX?: A Mesh is how JAX names the accelerators it will spread work across. Its docstring is blunt about the job: a Mesh "declares the hardware resources available in the scope of this manager." Concretely it is "a multidimensional NumPy array of JAX devices, where each axis of the mesh has a name, e.g. 'x' or 'y'." So you take a flat list of eight devices, reshape it into a 2×4 grid, and name the two axes — `Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))`. Nothing is sharded yet; the mesh is only the coordinate system. Everything else refers to devices by these axis names rather than by hardware id, which is what lets the same program run on a different-shaped cluster by rebuilding just the mesh.
What is a PartitionSpec in JAX?: A PartitionSpec says which array dimension is split across which mesh axis. Its docstring: it is a "Tuple describing how to partition an array across a mesh of devices," where "each element is either None, a string, or a tuple of strings." Position matters — the i-th element of the spec describes the i-th dimension of the array. `PartitionSpec('x', 'y')` "says that the first dimension of data is sharded across x axis of the mesh, and the second dimension is sharded across y axis of the mesh." A `None` leaves that dimension replicated (every device along that axis keeps a full copy), and a tuple like `('x', 'y')` shards one array dimension across two mesh axes at once.
What is NamedSharding in JAX?: NamedSharding is the pairing that turns a mesh and a spec into an actual placement. As the docstring puts it, a NamedSharding "expresses sharding using named axes" and "is a pair of a Mesh of devices and PartitionSpec which describes how to shard an array across that mesh." You build one as `NamedSharding(mesh, P('x', 'y'))` and hand it to `jax.device_put` to lay a real array onto the devices, or attach it to a jit input or output. Because it is expressed in mesh-axis names rather than device ids, the same NamedSharding is meaningful on any mesh with the same axis names.
How does jit decide the sharding of intermediate values?: By propagation. You annotate the shardings you care about — some inputs, sometimes an intermediate pinned with `jax.lax.with_sharding_constraint` — and inside jit the partitioner (GSPMD, and now Shardy) propagates those constraints through the rest of the program, filling in every unspecified value with a sharding it infers from the ops. When a computation needs data that lives on other devices, the same pass inserts the communication (an all-reduce over a summed-away axis, an all-gather to un-shard, a reduce-scatter). The catch is that propagation is heuristic: if "there aren't enough sharding constraints in the program, so a small change in propagation order can affect the final result," you can get a surprising layout or an out-of-memory. The fix is to pin the shardings that matter rather than hope the compiler guesses them.
What is the difference between shard_map and jit auto-sharding?: They sit at opposite ends of who writes the communication. Under jit auto-sharding you program with global arrays, annotate what you can, and let the compiler choose the rest and drop in the collectives — easy and composable, but the compiler's choices can surprise you. shard_map hands control back: it "maps a function over shards of data using a mesh of devices," running your function on each device against its local shard, and "the names of the Mesh can be used in collective communication operations in f" — so you call the all-reduce yourself. You specify in_specs and out_specs the same way, but there is no propagation and no compiler-inserted communication; the per-device shapes and the collectives are explicit. Auto for productivity, shard_map for predictable performance and hand-tuned kernels.
What is Shardy and when does it replace GSPMD in JAX?: Shardy is JAX's next-generation partitioner, "a new partitioning system co-developed by GDM Model Scaling (author of PartIR) and XLA/CoreML teams (author of GSPMD)," which "aims to provide better usability and control to users, and will gradually replace GSPMD and PartIR." Per the migration doc, "after the migration is complete in March 2026, Shardy will be the only partitioner in JAX." You can tell it is on by the log line "Using Shardy for XLA SPMD propagation," and until the cutover it can be turned off with `jax.config.update('jax_use_shardy_partitioner', False)`. Because Shardy uses a different propagation order and conflict-resolution heuristics, it "can sometimes output slightly different results" than GSPMD — which is usually a sign the program was under-constrained rather than that either partitioner is wrong.

Sharding in JAX

Questions

Related explainers