a mastery course · the xla path · through the compiler
Learn how XLA works.
The JAX path treats the compiler as a black box that keeps its promises. This path opens it: the seam every framework enters through, the IR where two hundred passes work, the floors where backends part ways, and the runtime stack above that decides who controls the fleet.
your progress 0/14 chapters
From the PJRT seam to the codegen floors: how a StableHLO program becomes an executable, with every decision visible in a dump.
- ch 01 PJRT Every frontend that talks to XLA, JAX included, crosses the exact same runtime seam to get there.
- ch 02 Ingestion What jit hands XLA and what XLA actually optimizes are two different IRs, and confusing them is the easiest way to misread a dump.
- ch 03 HLO By the time a program reaches HLO, every optimization decision XLA will ever make happens on this one representation and nowhere else.
- ch 04 Pipeline Run one program through a real compile and the dump folder holds every decision the compiler made, in order, as dozens of numbered files.
- ch 05 Fusion XLA has exactly one lever on memory traffic, and it never touches the math.
- ch 06 Layout and memory Two arrays with the same shape can still disagree about where every number actually sits in memory.
- ch 07 SPMD The mesh and the PartitionSpec you wrote from the JAX path did not vanish; they became the input to a pass that rewrites the whole module around them.
- ch 08 Collectives A collective instruction looks like one op in the dump; underneath, every device in the clique has to show up and call it in the same order or the whole fleet stalls.
- ch 09 The floors One fusion, dumped on three chips, comes back three different kinds of code.
- ch 10 Measure when models run out When XLA cannot decide which kernel wins on paper, it stops guessing and runs a race.
IFRT, multi-controller JAX, single-controller Pathways, and a capstone where the dump is the referee.
- ch 11 The array moves up PJRT hands a caller a pile of per-device buffers. IFRT hands back one array.
- ch 12 Multi-controller JAX Stock JAX has no coordinator: every host in the job runs the identical script and simply trusts that every other host is doing the same.
- ch 13 Single controller One Python program drives an entire pod of chips, and the workers underneath it never see your script at all.
- ch 14 The dump is the referee Three projects, one deliverable each, and one rule that applies to all three: no claim survives without a dump or a measurement standing behind it.
Facts are anchored to the public source tree at a verified commit, and readings are public only: openxla.org, the GitHub tree, the GSPMD and Pathways papers. Every cited source path exists; every reading URL answered when this path was built.
The pass pipeline shown is a real dump from a real compile on this machine (jax 0.4.38, CPU, 2026-07-27), and the HLO excerpts are the compiler's own output with paths shortened. Closed layers stay marked closed: libtpu below, the Pathways runtime above.
The JAX path teaches the language that feeds this compiler; the kernel path lives on the floor where XLA stops. The three cross-reference, and the xla gym floor drills what these chapters teach.