the xla path · 0/15
start the path

the xla path · Ingestion · lesson 01 of 1

The translation walk

One six-line function, three representations: the Python you wrote, the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → tracing recorded, and the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → received. This lesson walks the same program across both borders.

the goal Given any traced program, predict where each piece of it lands in the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →: which equation becomes which op, where the carry goes, and what disappears.

mastery work · this chapter0/2
  2. go →
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
§ 01

The program, on the source side

The walk needs a program small enough to hold whole and rich enough to exercise the interesting machinery. A running mean over rows does it: there is a carry (the total and the count), a loop (lax.scan), and elementwise arithmetic inside the step. Six lines of Python, and every one of them will be findable on the far side of two translations.

the program, verbatim from the corpus (generated by gen_cf_corpus.py, jax 0.4.38)
def running_mean(xs):
    def step(carry, x):
        total, n = carry
        return (total + x, n + 1.0), (total + x) / (n + 1.0)
    (_, _), means = jax.lax.scan(step, (jnp.zeros(xs.shape[1]), 0.0), xs)
    return means
§ 02

The jaxpr side

Trace it and the recording shows the grammar the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → lessons taught: one scan equation at the top level, with the step function as a nested jaxpr inside its params and the carry as explicit state. Notice what already happened here, before XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → saw anything: the tuple you wrote became positional values, the closure over nothing stayed empty, and every value carries its shape.

the head of the traced jaxpr, verbatim; the whole dump unfolds below
{ lambda ; a:f32[8,16]. let
    b:f32[16] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(16,)
      sharding=None
    ] 0.0
    _:f32[16] _:f32[] c:f32[8,16] = scan[
      _split_transpose=False
      jaxpr={ lambda ; d:f32[16] e:f32[] f:f32[16]. let
          g:f32[16] = add d f
          h:f32[] = add e 1.0
          i:f32[16] = add d f
          j:f32[] = add e 1.0
          k:f32[] = convert_element_type[new_dtype=float32 weak_type=False] j
the full jaxpr · 24 lines 
{ lambda ; a:f32[8,16]. let
    b:f32[16] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(16,)
      sharding=None
    ] 0.0
    _:f32[16] _:f32[] c:f32[8,16] = scan[
      _split_transpose=False
      jaxpr={ lambda ; d:f32[16] e:f32[] f:f32[16]. let
          g:f32[16] = add d f
          h:f32[] = add e 1.0
          i:f32[16] = add d f
          j:f32[] = add e 1.0
          k:f32[] = convert_element_type[new_dtype=float32 weak_type=False] j
          l:f32[16] = div i k
        in (g, h, l) }
      length=8
      linear=(False, False, False)
      num_carry=2
      num_consts=0
      reverse=False
      unroll=1
    ] b 0.0 a
  in (c,) }
§ 03

The StableHLO side

Lower it and the same structure re-appears in the other notation. The scan equation became stablehlo.while; the carry you threaded became the operand tuple riding the loop, visible in the %iterArg list; the nested jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → became the cond and do regions attached to the op. Nothing about the computation moved, and everything about the notation did.

the while op with its carry tuple, verbatim from the same corpus entry
    %2:5 = stablehlo.while(%iterArg = %arg0, %iterArg_2 = %c, %iterArg_3 = %0, %iterArg_4 = %cst_0, %iterArg_5 = %1) : tensor<8x16xf32>, tensor<i32>, tensor<16xf32>, tensor<f32>, tensor<8x16xf32>
     cond {
      %c_6 = stablehlo.constant dense<8> : tensor<i32>
      %3 = stablehlo.compare  LT, %iterArg_2, %c_6,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>
      stablehlo.return %3 : tensor<i1>
    } do {
the full stablehlo module · 48 lines 
module @jit_running_mean attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<8x16xf32>) -> (tensor<8x16xf32> {jax.result_info = ""}) {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
    %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %cst_1 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %1 = stablehlo.broadcast_in_dim %cst_1, dims = [] : (tensor<f32>) -> tensor<8x16xf32>
    %c = stablehlo.constant dense<0> : tensor<i32>
    %2:5 = stablehlo.while(%iterArg = %arg0, %iterArg_2 = %c, %iterArg_3 = %0, %iterArg_4 = %cst_0, %iterArg_5 = %1) : tensor<8x16xf32>, tensor<i32>, tensor<16xf32>, tensor<f32>, tensor<8x16xf32>
     cond {
      %c_6 = stablehlo.constant dense<8> : tensor<i32>
      %3 = stablehlo.compare  LT, %iterArg_2, %c_6,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>
      stablehlo.return %3 : tensor<i1>
    } do {
      %c_6 = stablehlo.constant dense<0> : tensor<i32>
      %3 = stablehlo.compare  LT, %iterArg_2, %c_6,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>
      %4 = stablehlo.convert %iterArg_2 : tensor<i32>
      %c_7 = stablehlo.constant dense<8> : tensor<i32>
      %5 = stablehlo.add %4, %c_7 : tensor<i32>
      %6 = stablehlo.select %3, %5, %iterArg_2 : tensor<i1>, tensor<i32>
      %c_8 = stablehlo.constant dense<0> : tensor<i32>
      %7 = stablehlo.dynamic_slice %iterArg, %6, %c_8, sizes = [1, 16] : (tensor<8x16xf32>, tensor<i32>, tensor<i32>) -> tensor<1x16xf32>
      %8 = stablehlo.reshape %7 : (tensor<1x16xf32>) -> tensor<16xf32>
      %9:3 = func.call @None(%iterArg_3, %iterArg_4, %8) : (tensor<16xf32>, tensor<f32>, tensor<16xf32>) -> (tensor<16xf32>, tensor<f32>, tensor<16xf32>)
      %10 = stablehlo.broadcast_in_dim %9#2, dims = [1] : (tensor<16xf32>) -> tensor<1x16xf32>
      %11 = stablehlo.compare  LT, %iterArg_2, %c_6,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>
      %12 = stablehlo.convert %iterArg_2 : tensor<i32>
      %13 = stablehlo.add %12, %c_7 : tensor<i32>
      %14 = stablehlo.select %11, %13, %iterArg_2 : tensor<i1>, tensor<i32>
      %15 = stablehlo.dynamic_update_slice %iterArg_5, %10, %14, %c_8 : (tensor<8x16xf32>, tensor<1x16xf32>, tensor<i32>, tensor<i32>) -> tensor<8x16xf32>
      %c_9 = stablehlo.constant dense<1> : tensor<i32>
      %16 = stablehlo.add %iterArg_2, %c_9 : tensor<i32>
      stablehlo.return %iterArg, %16, %9#0, %9#1, %15 : tensor<8x16xf32>, tensor<i32>, tensor<16xf32>, tensor<f32>, tensor<8x16xf32>
    }
    return %2#4 : tensor<8x16xf32>
  }
  func.func private @None(%arg0: tensor<16xf32>, %arg1: tensor<f32>, %arg2: tensor<16xf32>) -> (tensor<16xf32>, tensor<f32>, tensor<16xf32>) {
    %0 = stablehlo.add %arg0, %arg2 : tensor<16xf32>
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
    %1 = stablehlo.add %arg1, %cst : tensor<f32>
    %2 = stablehlo.add %arg0, %arg2 : tensor<16xf32>
    %3 = stablehlo.add %arg1, %cst : tensor<f32>
    %4 = stablehlo.convert %3 : tensor<f32>
    %5 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<f32>) -> tensor<16xf32>
    %6 = stablehlo.divide %2, %5 : tensor<16xf32>
    return %0, %1, %6 : tensor<16xf32>, tensor<f32>, tensor<16xf32>
  }
}
§ 04

The translator's dictionary

The walk generalizes to a small dictionary you can apply to any program. Equations become ops one for one, with the primitive name usually surviving recognizably. Nested jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → become regions on the op that carried them. Explicit carry state stays explicit, changing only its clothes: params-and-binders on one side, an operand tuple on the other. And names die at the border: your Python identifiers were already gone in the jaxpr, and the jaxpr's letter variables give way to numbered % values. What survives every translation is shapes, dtypes, and structure, which is exactly the set of things the compiler is allowed to care about.

The chapter above this lesson explains why the two IRs exist at all; the corpus x-ray on the gym floor has seventeen more programs to walk. Fluency is the point: the border should feel like notation, not like a wall.

before you move on

Check yourself

01 Where does the scan's carry live on each side of the translation?

In the jaxpr, as explicit state on the scan equation with the body nested in its params. In StableHLO, as the operand tuple threading through the while op's regions. Same contract, two notations.

02 What survives translation unchanged, and what changes form?

Shapes, dtypes, and the structure of the computation survive. Notation changes: equations become ops, nested jaxprs become regions, and names die at each border, ending as numbered values.

assigned

Readings