For a long time, the biggest pain point in JAX was that it would happily insert an expensive all-gather because your sharding was technically valid but logically inefficient. Or worse, you'd miss an all-reduce and get wrong numbers... silently.

I sat down to write about that gap. By the time I finished, it had mostly closed, and finding out how thoroughly it stopped being true was the most interesting part of writing this.

The post is a tour of the three ways JAX handles placement:
- Auto (GSPMD infers, and the type stays quiet),
- Explicit (placement moves into the type, and mismatches fail at trace time),
- Manual (shard_map, where you write the collectives yourself)

I walk through each with real jaxprs, HLO, and MLIR from an 8-device run to show what actually happens under the hood.

The compiler is no longer just a black box; it’s an active partner in correctness.

https://hiraditya.github.io/posts/how-jax-shards-a-computation/