Six shardings, one matmul

The corpus behind the collective picker instrument (EX·16 on the spmd chapter) holds six real cases: one matmul, shapes (512,256) x (256,512), on an eight-device mesh laid out (4 data, 2 model), jax 0.4.38. Each row names the two input shardings, the requested output sharding, and the collectives the compiled module actually contains. The table is worth reading before any theory, because two of its rows are silent and one of them is louder than you would guess.

Rows one and four insert nothing, and they share a shape: the contraction axis arrives whole on every shard, and the local results already sit exactly as the requested output sharding wants them. Row one splits a by rows across data and leaves b replicated, so each shard computes its own rows of the answer; row four mirrors it on the other side, splitting b by columns. No resharding question is ever asked, so the partitioner has nothing to answer. Free is not a special case; it is what alignment looks like.

Row three is the teaching row. Both inputs are split along the contraction axis, a over model and b over data, so no shard can finish its arithmetic alone; something has to move. A partial-sum analysis predicts a reduction. The capture shows an all-gather: the partitioner chose to reassemble an operand and compute locally rather than reduce partials, and the module records that choice without asking your opinion. Row six goes further, splitting both inputs over the same model axis, and pays with two collectives, an all-gather and a collective-permute, to reconcile the two different splits.

The lesson in those rows is not the specific plans; it is who decides. The partitioner picks among mathematically equal strategies with a cost model you do not see, and the only reliable account of the plan you got is the dump. The partitioner lesson under the spmd unit gives every collective its author; this table is where you watch the author make six different calls on one program. Step through the same six rows live in EX·16, then predict row five before revealing it.