Ingestion

Every compile starts the same way: a frontend emits StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, hands it to PJRT, and PJRT converts it into something else before a single optimization pass runs. StableHLO is a versioned, portable dialect of MLIR, vendored inside 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 tree so every frontend can rely on the exact same wire format regardless of which compiler version is on the other end. It is designed to be read, written, and serialized across process and even version boundaries.

That conversion is not cosmetic. The function doing it, ConvertStablehloToHlo (xla/hlo/translate/stablehlo.h), with the PJRT-side wrapper sitting at xla/pjrt/mlir_to_hlo.cc, turns a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → module into HLO, 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 internal representation, and from that point on StableHLO is gone. Nothing downstream of ingestion ever looks at it again.

StableHLO is the wire; HLO is the workbench.

The natural assumption, especially if you have spent any time in the MLIR ecosystem, is that HLO must be another MLIR dialect sitting one step below StableHLO. It is not. HLO is a hand-rolled, mutable C++ graph IR (xla/hlo/ir/hlo_module.h, xla/hlo/ir/hlo_instruction.h) that predates MLIR entirely, with its own class hierarchy, its own mutation APIs, and none of MLIR's operation or attribute machinery underneath it.

There is a third IR in the picture that makes the confusion worse: MHLO, 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 legacy MLIR dialect, still lives in the tree at xla/mlir_hlo/. PJRT does not speak it. Round-trip utilities between MHLO and HLO exist purely for older tooling that has not migrated yet; on the path this chapter is teaching, from a jit call to a running executable, MHLO never appears.

The layering earns its keep rather than being an accident of history. StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →'s whole job is to be stable: it carries explicit versioning and compatibility guarantees, so a program serialized by one compiler version can still be read by a slightly different one on the other end of a network or a saved model. HLO has no interest in any of that. Its job is to be mutated, by pass after pass, in place, and it is built for exactly that: XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → runs upward of two hundred passes against it, each one free to rewrite the graph directly.

One format optimized for staying the same across versions, one format optimized for changing constantly during a single compile. Trying to make a single IR do both jobs well is what the split avoids.

This is what jit actually hands to XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →, captured from a small attention program before any optimization pass has touched it. Read the shape annotations closely: every tensor states its dtype and dimensions inline, every op names its exact operands, and the contraction in the dot_generalThe general matmul primitive. Its dimension_numbers state which axes contract and which batch; decoding them on sight is the core IR-reading skill.taught in /l/stablehlo → spells out which axes multiply against which, information Python never made you write down explicitly.

module @jit_attend attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<64x64xf32>, %arg1: tensor<64x64xf32>, %arg2: tensor<64x64xf32>) -> (tensor<64x64xf32> {jax.result_info = ""}) {
    %0 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<64x64xf32>) -> tensor<64x64xf32>
    %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<64x64xf32>, tensor<64x64xf32>) -> tensor<64x64xf32>
    %cst = stablehlo.constant dense<6.400000e+01> : tensor<f32>
    %2 = stablehlo.sqrt %cst : tensor<f32>
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<f32>) -> tensor<64x64xf32>
    %4 = stablehlo.divide %1, %3 : tensor<64x64xf32>