Two dumps, two truths

Every Pallas or JAX kernel you write for TPU passes through XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → twice in ways that matter to you directly: once when it gets translated, and once when it gets decided. Both moments produce a dump you can read as plain text, but the two dumps answer different questions, and mixing them up is one of the fastest ways to draw the wrong conclusion about why a kernel runs slow.

jax.jit(fn).lower(*args).as_text() hands you StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, the program almost exactly as you wrote it, translated into XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s portable intermediate representation but not yet touched by any backend. Shapes are fixed, but nothing about how the TPU will actually run this has been decided. No fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →, no layout, no buffer assignment. It reads like source code with the syntax swapped, not like an execution plan.

jax.jit(fn).compile(*args).as_text() is the other call, and it hands you something the backend actually produced. Elementwise ops have merged into fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → ops. Tensors have picked a physical layout. Buffers have been assigned to specific memory. This dump is backend specific, since running the same function through a different backend produces a different structure, because the decisions inside it are compiler decisions made for one target. When you want to know why a kernel is slow, this is the only dump that can answer, because it is the only one reflecting what actually ran.

The mistake worth naming early is reading the lowered StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and reasoning about performance from it. Nothing in that dump has been fused, tiled, or laid out yet, so a loop that looks expensive there might vanish into a single fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → once compiled, and a loop that looks cheap might turn into three separate passes over HBM. Performance questions belong to the compiled dump, always. The lowered dump answers a different question: did you write the program you meant to write, before the compiler gets a say in any of it.

In practice, every performance investigation should start the same way: call .compile(*args).as_text(), not .lower(*args).as_text(), and look for the shapes and fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → kinds that actually made it into the final program. The lowered dump still has its place, mostly for confirming that the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → your Python traced to matches what you intended, before any compiler decision has had the chance to obscure it.

The size difference between the two dumps is itself a signal. A .lower() dump for a modest kernel might run a few dozen lines, mirroring your Python almost one to one. The .compile() dump for the same kernel can run many times longer, because every operand now carries a layout annotation, every fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → group is spelled out with its inputs and outputs, and buffer assignment adds allocation metadata that had no reason to exist before any backend decision was made.