What changes from the jaxpr
StableHLO is the 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. 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.