Reading nested jaxprs

Some primitives don't just carry numbers or axes in their params. They carry entire jaxprs. lax.cond carries its branches this way, one nested jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → per branch, each a complete function from the same inputs to the same output shape, and the primitive picks which one actually runs at execution time. lax.scan carries its loop body the same way, a single nested jaxpr applied once per iteration with the carry threaded in and out. In both cases, the params in square brackets that held a shape or a dtype for a simpler primitive now hold something built from the full grammar of the first section.

Printed output shows this nesting through indentation. A scan or cond equation appears at the outer level like any other, and beneath it, indented, sits the body it carries: its own invars, its own equations, its own outvars. Read that inner block exactly the way you read the outer one, because it is built from the same grammar, variables bound once, params in brackets, an outvars line at the end. The only thing that changes as you go deeper is how far it sits indented on the page. Nothing about how to read it changes.

{ 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,) }

This matters past the classroom example, because kernel jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → nest the same way. A pipelined Pallas kernel has a grid, and the grid's per-step body is exactly the kind of nested jaxpr that scan and cond produce here, an inner function with its own bindings that gets run once per step by the compiler. If you can read the scan example in this chapter, you can read that kernel jaxpr too, because it's the identical grammar one level deeper. Getting comfortable with nesting now is what keeps the kernel material legible later, instead of reading as a wall of unfamiliar syntax.