This is what leaves JAX and crosses the seam, the form every chapter at the waist works on.
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%1 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%2 = stablehlo.dot_general %arg0, %arg3, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%3 = stablehlo.transpose %1, dims = [1, 0] : (tensor<16x32xf32>) -> tensor<32x16xf32>
%4 = stablehlo.dot_general %0, %3, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x16xf32>) -> tensor<16x16xf32>
%cst = stablehlo.constant dense<3.200000e+01> : tensor<f32>
%5 = stablehlo.sqrt %cst : tensor<f32>
%6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<f32>) -> tensor<16x16xf32>
%7 = stablehlo.divide %4, %6 : tensor<16x16xf32>What changes from the jaxpr
StableHLO is the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →'s information re-expressed in MLIR form, versioned and stable so that different tools and framework versions can exchange it. Same ops, same shapes, now written as a module with a function, in the syntax the whole compiler ecosystem shares. Print it with jax.jit(fn).lower(*args).as_text().
The word MLIR will follow you from here down, so pin it now: MLIR is not a language, it is a framework for defining IRs (called dialects) and transformations over them. StableHLO is a dialect. MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic →, two layers down, is a dialect. When you read that something is "an MLIR pass," it means one bounded rewrite applied to one of these dialects.
func.func public @main(%arg0: tensor<128x64xbf16>, %arg1: tensor<128x64xbf16>, %arg2: tensor<128x64xbf16>) -> (tensor<128x64xbf16> {jax.result_info = ""}) {
%0 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<128x64xbf16>) -> tensor<64x128xbf16>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<128x64xbf16>, tensor<64x128xbf16>) -> tensor<128x128xbf16>
%cst = stablehlo.constant dense<0xFF80> : tensor<bf16>
%2 = stablehlo.reduce(%1 init: %cst) applies stablehlo.maximum across dimensions = [1] : (tensor<128x128xbf16>, tensor<bf16>) -> tensor<128xbf16>
%3 = stablehlo.broadcast_in_dim %2, dims = [0] : (tensor<128xbf16>) -> tensor<128x1xbf16>
%4 = stablehlo.broadcast_in_dim %3, dims = [0, 1] : (tensor<128x1xbf16>) -> tensor<128x128xbf16>
%5 = stablehlo.subtract %1, %4 : tensor<128x128xbf16>
%6 = stablehlo.exponential %5 : tensor<128x128xbf16>The fluency drill
One skill separates people who can read tensor IRs from people who squint at them: decoding dot_general dimension_numbers on sight. Which axes contract, which axes batch. Drill it here where the notation is friendly, because every matmul you ever chase through a compiler dump is one of these.
A discipline this layer rewards: when something surprises you in a StableHLO dump, resist the urge to guess. The dump is complete; the answer is in it. The EX·05 instrument on Stage 2 holds this exact program open at three layers with hover sync, which makes the correspondence physical.
Lessons
- 01The type systemEvery value carries its type in the text: dimensions first, element type second, and a scalar is just a tensor with no dimensions. ·
- 02Control flow is regionsNo gotos, no basic blocks: a while or an if is one op holding nested bodies, and the op decides which body runs. ·
- 03custom_call, the escape hatchThe opset is closed on purpose, and everything outside it, your Pallas kernels included, travels through one op. ·
- 04Why it is called StableThe name is a serialization promise with numbers attached, and it is the reason a saved model outlives the toolchain that made it. ·
- 05The HLO family treeFive names one letter apart, and none of them is a rung on the same ladder. Each one exists because something broke that the one before it could not fix. ·