the path · 0/15
start the path

the kernel path · StableHLO · lesson 02 of 5

Control flow is regions

No gotos, no basic blocks: a while or an if is one op holding nested bodies, and the op decides which body runs.

the goal Find your scan’s carry and your cond’s branches inside a real dump, and read region nesting as ordinary control flow.

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

Control flow is regions

StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → has no goto and no basic-block jump the way LLVM IR does. Control flow shows up as regions, nested bodies attached to a single op, and the op decides which region runs and how results flow back out. That single design choice is why a lax.scan and a lax.cond in your JAX code become such literal, readable structures once they hit StableHLO. The loop or branch is not scattered across labels, it is a self-contained block sitting right where the op is called.

control flow as regions: the carry is explicit loop state, and both regions live inside the op
init carry stablehlo.while · carry: (total, n) cond region n < 8 → i1 body region new carry = step(carry) true carry loops back false: final carry lax.scan lowers to exactly this, with the stacked outputs threaded through the carry

lax.scan lowers to stablehlo.while, and the carry you thread through your scan body becomes explicit loop state on the op itself, not a hidden closure. The while carries a tuple of operands into its body region on every iteration and receives an updated tuple back out, and a separate condition region decides whether to keep going. Nothing about the carry is implicit. If your scan carries a running sum and an index, both show up as named operands on stablehlo.while, updated by name each pass, so you can trace exactly which value changed on which iteration.

lax.scan becomes stablehlo.while · generated by gen_cf_corpus.py
  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>
  }
}
the full stablehlo · 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>
  }
}

lax.cond lowers to stablehlo.if, and both branches are present in the dump at once, as two separate regions attached to the same op, not as a single path chosen at trace time. The predicate, a tensor<i1>, selects which region actually executes; the other region still sits there in the text, fully typed, ready to run if the predicate flips. That is different from a Python if statement, which only ever traces one branch. StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → keeps both, because the compiler has to reason about the program for every possible input, not just the one you traced with.

lax.cond becomes stablehlo.if · generated by gen_cf_corpus.py
  func.func public @main(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<f32>) -> (tensor<16xf32> {jax.result_info = ""}) {
    %0 = call @norm(%arg1) : (tensor<16xf32>) -> tensor<f32>
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
    %1 = stablehlo.compare  GT, %0, %cst,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1>
    %2 = stablehlo.convert %1 : (tensor<i1>) -> tensor<i32>
    %3 = "stablehlo.case"(%2) ({
      %4 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<16xf32>
      %5 = stablehlo.multiply %4, %arg1 : tensor<16xf32>
      %6 = stablehlo.subtract %arg0, %5 : tensor<16xf32>
      stablehlo.return %6 : tensor<16xf32>
    }, {
      %4 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<16xf32>
      %5 = stablehlo.multiply %4, %arg1 : tensor<16xf32>
      %6 = func.call @norm(%arg1) : (tensor<16xf32>) -> tensor<f32>
      %7 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor<f32>) -> tensor<16xf32>
      %8 = stablehlo.divide %5, %7 : tensor<16xf32>
      %9 = stablehlo.subtract %arg0, %8 : tensor<16xf32>
      stablehlo.return %9 : tensor<16xf32>
    }) : (tensor<i32>) -> tensor<16xf32>
    return %3 : tensor<16xf32>
  }
  func.func private @norm(%arg0: tensor<16xf32>) -> tensor<f32> {
    %0 = stablehlo.multiply %arg0, %arg0 : tensor<16xf32>
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %1 = stablehlo.reduce(%0 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
    %2 = stablehlo.sqrt %1 : tensor<f32>
    return %2 : tensor<f32>
  }
}
the full stablehlo · 30 lines 
module @jit_clipped_update attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<f32>) -> (tensor<16xf32> {jax.result_info = ""}) {
    %0 = call @norm(%arg1) : (tensor<16xf32>) -> tensor<f32>
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
    %1 = stablehlo.compare  GT, %0, %cst,  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1>
    %2 = stablehlo.convert %1 : (tensor<i1>) -> tensor<i32>
    %3 = "stablehlo.case"(%2) ({
      %4 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<16xf32>
      %5 = stablehlo.multiply %4, %arg1 : tensor<16xf32>
      %6 = stablehlo.subtract %arg0, %5 : tensor<16xf32>
      stablehlo.return %6 : tensor<16xf32>
    }, {
      %4 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<16xf32>
      %5 = stablehlo.multiply %4, %arg1 : tensor<16xf32>
      %6 = func.call @norm(%arg1) : (tensor<16xf32>) -> tensor<f32>
      %7 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor<f32>) -> tensor<16xf32>
      %8 = stablehlo.divide %5, %7 : tensor<16xf32>
      %9 = stablehlo.subtract %arg0, %8 : tensor<16xf32>
      stablehlo.return %9 : tensor<16xf32>
    }) : (tensor<i32>) -> tensor<16xf32>
    return %3 : tensor<16xf32>
  }
  func.func private @norm(%arg0: tensor<16xf32>) -> tensor<f32> {
    %0 = stablehlo.multiply %arg0, %arg0 : tensor<16xf32>
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %1 = stablehlo.reduce(%0 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
    %2 = stablehlo.sqrt %1 : tensor<f32>
    return %2 : tensor<f32>
  }
}

Reading these regions gets easier with practice, and the fastest way to build that muscle is staring at real dumps until the nesting stops feeling foreign and starts feeling like ordinary control flow again. The corpus is full of scans and conds that started life as plain Python loops and branches, and once you have traced a handful of them by hand, spotting the carry inside a while region or the predicate driving an if stops taking any real effort.

before you move on

Check yourself

01 Where did your lax.scan’s carry go in the StableHLO dump?

Onto the stablehlo.while op itself, as an explicit tuple of operands carried into the body region every iteration and returned updated, not as a hidden closure.

02 After tracing chose one path for a Python if, why does a lax.cond dump still show two bodies?

cond lowers to stablehlo.if with both branches present as regions on the op; a tensor<i1> predicate selects which region executes at runtime, so the choice survived tracing instead of being resolved by it.

assigned

Readings