SPMD

The chapter above introduces this as the mechanism under JAX's Mesh and PartitionSpec, and names SpmdPartitioner correctly. Two details make the source navigable from there. The pass is an HloModulePass whose name() returns spmd-partitioning, so that is the string to grep for in a dump filename and the string to pass to a disable flag. And it is the pass that inserts all-gather, all-reduce, collective-permute, and all-to-all, which means every collective in your optimized module has exactly one author.

The design has a paper, and the paper's numbers are worth carrying because they set the scale this machinery was built for. GSPMD: General and Scalable Parallelization for ML Computation Graphs, by Xu, Lee, Chen, Hechtman and thirteen others, reports 50% to 62% compute utilization on up to 2048 Cloud TPUv3 cores for models with up to one trillion parameters. The design brief in the abstract is the part that explains the annotation model: practitioners write code as if targeting a single device, then distribute it through minimal annotations.

Minimal is the operative word, and it is a claim about the propagation algorithm rather than about ergonomics. The system determines operator partitioning from limited user guidance, which is to say most instructions in a real partitioned module carry a sharding nobody wrote.

Propagation runs first and completes before any rewriting starts. Annotations you gave flow forward and backward through the graph until nothing changes, and every unannotated instruction inherits whatever its neighbours settled on. The interesting cases are the disagreements, where two neighbours would prefer different shardings for the same value.

That those conflicts are ordinary rather than exceptional is visible in the options struct. SpmdPartitionerOptions carries a field called need_resolve_conflicts, which only exists because conflicting shardings are a routine input to this pass rather than a malformed one. Reading an options struct is often the fastest way to learn what a pass considers a normal day.

Where the annotation lives is currently moving. Shardy, the sdy dialect, attaches sharding in MLIR to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → before the module ever becomes HLO. That is an earlier altitude for the same information, not a different partitioner: the destination is still a module where every instruction has a sharding and the rewriting can begin.

The rewriting half has two supporting classes, and knowing them changes how you read a partitioned dump. SpmdBuilder wraps HloComputation::Builder and tracks derived instructions and broadcast dimensions as it goes, so the partitioner can reason about what it just created. PartitionedHlo is the more useful one to know: it represents a value together with its current partitioned state, and it exposes the operations that change that state, including Reshard(), PadWithZero(), and Replicate().

So a collective in your module is the output of a resharding request. Something asked for a value in a sharding it did not currently have, PartitionedHlo produced the instructions that get it there, and the collective you are looking at is that answer. Reading a dump backwards from a collective becomes a specific question with a specific answer: which two shardings does this instruction sit between, and which consumer demanded the second one.

PartitionedHlo also caches its reshards, which is the reason the same conversion appearing twice in your Python does not always appear twice in the module. If you are counting collectives to reason about network time, count them in the optimized HLO rather than in your source, and expect the two counts to differ.

The partitioner decides that an all-reduce belongs at a particular point in the dataflow. It does not decide how that all-reduce runs. An HLO collective names the operation, the participant groups, and the shapes; it says nothing about ring versus tree, about which links carry which fragment, or about how the transfer is chunked.

That choice belongs to the runtime, below the compiler and below the HLO you are reading. On TPU it lives inside libtpu, whose own description names inter-chip communication alongside compilation and runtime execution as the three things the library provides. The 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 → topology the kernel path teaches is a property of the physical pod, and the collective implementation is chosen against it at a level the dump does not show you.

There is one compiler-visible handle on that boundary, and /xla/collectives is where the path picks it up: the async pair. A collective split into all-reduce-start and all-reduce-done gives the scheduler a gap to fill with independent compute, and the width of that gap is a compiler decision made in full ignorance of which link the bytes will cross. Both layers optimize the same collective and neither can see the other's choice.

Propagation is not the only way to get a partitioned program, and knowing where the other door is matters for reading dumps. shard_map is JAX's manual mode, described in its own documentation as a single-program multiple-data multi-device parallelism API to map a function over shards of data. The contrast with jit is stated directly in the same document: with jit you write code as if for a single device and the compiler partitions it and generates the collectives behind the scenes.

Inside a shard_map you do neither of those things. The function body sees per-device shapes already, you write the communication yourself, and the available primitives are a short list: psum for an all-reduce sum, all_gather, all_to_all, psum_scatter for reduce-scatter, and axis_index when a device needs to know which one it is. The mesh, in_specs, and out_specs arguments say how the inputs get split and how the outputs get reassembled.

The two modes compose rather than compete. The documented pattern is manual control across groups of devices with compiler-based partitioning inside each group, which is what pipelining across islands with data parallelism inside them looks like. When you read a dump from such a program, the manual regions are the parts where the shapes are already per-device and the partitioner had nothing to propagate, and the collectives there are yours rather than the pass's.

Inside a manual region, every collective in the dump is one you wrote.

# propagated: you annotate the edges, the partitioner fills in the middle
y = jax.jit(f, out_shardings=NamedSharding(mesh, P("data", None)))(x)


# manual: the body sees per-device shapes and you write the communication
@partial(shard_map, mesh=mesh, in_specs=P("data", None), out_specs=P(None, None))
def g(x_shard):
    local = x_shard @ w          # already a per-device shape here
    return jax.lax.psum(local, "data")   # the only collective in this region