the path · 0/15
start the path

the kernel path · jaxpr · lesson 02 of 5

Transforms are jaxpr rewrites

grad, vmap, scan, remat, shard_map: one small dense function, traced once, then rewritten five ways in front of you.

the goal Predict what each major transform does to a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → (which grow equations, which only move shapes, which nest bodies) before opening the gallery.

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

Transforms are jaxpr rewrites

The gallery function for this chapter is small on purpose: dense(x, w, b) = tanh(x @ w + b), traced at x shaped (32, 64), w shaped (64, 64), and b shaped (64,). Every variant below starts from that same traced jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → and asks what one specific transform does to it. The answer is never a rewrite of the Python source; JAX doesn't re-run your function through some separate path for grad versus vmap. It rewrites the jaxpr the first trace produced, which is why comparing the jaxprs side by side tells you exactly what each transform does.

Call jax.grad on dense and the forward jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → does not disappear, it grows a second half. Reverse-mode autodiff appends the backward pass as more equations in the same jaxpr: transpose equations undo the layout dot_general needed, mul and sub equations reconstruct the tanh derivative, 1 - tanh(x)^2, sitting right around the original tanh equation, and add_any equations accumulate cotangents wherever a value fed more than one downstream use. This isn't a separate function running alongside the original. It's the same jaxpr grammar, just with more equations in it.

under jax.grad · the forward is joined by its transpose: find tanh's derivative (1 - tanh^2) as mul equations · generated by gen_transform_gallery.py
{ lambda ; a:bf16[32,64] b:bf16[64,64] c:bf16[64]. let
    d:bf16[32,64] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] a b
    e:bf16[1,64] = broadcast_in_dim[
      broadcast_dimensions=(1,)
      shape=(1, 64)
      sharding=None
    ] c
    f:bf16[32,64] = add d e
    g:bf16[32,64] = tanh f
    h:bf16[32,64] = sub 1 g
    i:f32[32,64] = convert_element_type[new_dtype=float32 weak_type=False] g
    j:f32[] = reduce_sum[axes=(0, 1)] i
    _:bf16[] = convert_element_type[new_dtype=bfloat16 weak_type=False] j
    k:f32[32,64] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(32, 64)
      sharding=None
    ] 1.0
    l:bf16[32,64] = convert_element_type[new_dtype=bfloat16 weak_type=False] k
    m:bf16[32,64] = mul l h
    n:bf16[32,64] = mul m g
    o:bf16[32,64] = add_any m n
    p:bf16[64,64] = dot_general[
      dimension_numbers=(([0], [0]), ([], []))
      preferred_element_type=bfloat16
    ] o a
    q:bf16[64,64] = transpose[permutation=(1, 0)] p
  in (q,) }

Call jax.vmap instead and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → looks almost identical to the original trace. Every shape in the traced program gains a leading batch dimension of 8, matching the axis you mapped over, while the equation list itself stays unchanged: dot_general is still dot_general, tanh is still tanh. vmap doesn't introduce a loop or a new primitive to walk over the batch. It adds one dimension to every operand and lets the existing primitives handle it, and that stability in the op set is intentional, not incidental.

under jax.vmap · every equation gains a batch dimension; nothing else changes, which is the whole point of vmap · generated by gen_transform_gallery.py
{ lambda ; a:bf16[8,32,64] b:bf16[64,64] c:bf16[64]. let
    d:bf16[8,32,64] = dot_general[
      dimension_numbers=(([2], [0]), ([], []))
      preferred_element_type=bfloat16
    ] a b
    e:bf16[1,64] = broadcast_in_dim[
      broadcast_dimensions=(1,)
      shape=(1, 64)
      sharding=None
    ] c
    f:bf16[1,1,64] = broadcast_in_dim[
      broadcast_dimensions=(np.int64(1), np.int64(2))
      shape=(1, 1, 64)
      sharding=None
    ] e
    g:bf16[8,32,64] = add d f
    h:bf16[8,32,64] = tanh g
  in (h,) }

jax.lax.scan changes the picture more directly: the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → now has one scan equation, and the loop body lives as a nested jaxpr inside that equation's params rather than as a flat sequence of equations. The carry, whatever state threads from one iteration to the next, is wired through the scan equation's inputs and outputs explicitly. Reading a scanned jaxpr means reading the outer equation for how the carry enters and leaves, then reading the inner jaxpr, which follows the same grammar as any other, for what a single step actually does.

under lax.scan · the body becomes one nested jaxpr run length times: read the carry in, the carry out · generated by gen_transform_gallery.py
{ lambda ; a:bf16[32,64] b:bf16[4,64,64] c:bf16[64]. let
    d:bf16[32,64] = scan[
      _split_transpose=False
      jaxpr={ lambda ; e:bf16[64] f:bf16[32,64] g:bf16[64,64]. let
          h:bf16[32,64] = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=bfloat16
          ] f g
          i:bf16[1,64] = broadcast_in_dim[
            broadcast_dimensions=(1,)
            shape=(1, 64)
            sharding=None
          ] e
          j:bf16[32,64] = add h i
          k:bf16[32,64] = tanh j
        in (k,) }
      length=4
      linear=(False, False, False)
      num_carry=1
      num_consts=1
      reverse=False
      unroll=1
    ] c a b
  in (d,) }

jax.remat shows up only once you also take a gradient. Inside the grad trace, the recomputed forward pass sits as a nested jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → marked by a remat2 primitive, marking exactly the region that gets recomputed on the backward pass instead of stored from the forward pass. It nests the same way a scan body nests: the params hold a jaxpr instead of a shape or an axis list, and reading it means stepping into that nested block like any other.

under jax.checkpoint · a remat wrapper marks the body for recompute in the backward pass; the body itself is unchanged · generated by gen_transform_gallery.py
{ lambda ; a:bf16[32,64] b:bf16[64,64] c:bf16[64]. let
    d:bf16[32,64] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] a b
    e:bf16[1,64] = broadcast_in_dim[
      broadcast_dimensions=(1,)
      shape=(1, 64)
      sharding=None
    ] c
    f:bf16[32,64] = add d e
    g:bf16[32,64] = tanh f
    h:f32[32,64] = convert_element_type[new_dtype=float32 weak_type=False] g
    i:f32[] = reduce_sum[axes=(0, 1)] h
    _:bf16[] = convert_element_type[new_dtype=bfloat16 weak_type=False] i
    j:f32[32,64] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(32, 64)
      sharding=None
    ] 1.0
    k:bf16[32,64] = convert_element_type[new_dtype=bfloat16 weak_type=False] j
    l:bf16[64,64] = remat2[
      differentiated=True
      jaxpr={ lambda ; m:bf16[32,64] n:bf16[64,64] o:bf16[64] p:bf16[32,64]. let
          q:bf16[32,64] = dot_general[
            dimension_numbers=(([1], [0]), ([], []))
            preferred_element_type=bfloat16
          ] m n
          r:bf16[1,64] = broadcast_in_dim[
            broadcast_dimensions=(1,)
            shape=(1, 64)
            sharding=None
          ] o
          s:bf16[32,64] = add q r
          t:bf16[32,64] = tanh s
          u:bf16[32,64] = sub 1 t
          v:bf16[32,64] = mul p u
          w:bf16[32,64] = mul v t
          x:bf16[32,64] = add_any v w
          y:bf16[64,64] = dot_general[
            dimension_numbers=(([0], [0]), ([], []))
            preferred_element_type=bfloat16
          ] x m
          z:bf16[64,64] = transpose[permutation=(1, 0)] y
        in (z,) }
      policy=None
      prevent_cse=True
    ] a b c k
  in (l,) }

jax.shard_map rewrites the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → around per-shard shapes rather than global ones: every array in the traced program is sized as it exists on one shard, not as the full logical array. pbroadcast equations appear where a value needs to reach every shard, but in this traced program no collective communication primitive shows up at all, because dense doesn't need one. Every shard can compute its own tanh, matmul, and add independently, and the jaxpr tells you that directly, before you would ever profile it.

before you move on

Check yourself

01 Under jax.grad, does the forward jaxpr get replaced by a derivative program? What actually happens?

No. The forward equations stay and reverse-mode autodiff appends the backward pass as more equations in the same jaxpr: transposes undoing layouts, mul and sub chains computing the derivative, flowing the cotangent back.

02 You diff a jaxpr before and after vmap and the equation list is identical. Where did the transform go?

Into the shapes. Every traced value gained a leading batch dimension matching the mapped axis; dot_general is still dot_general, applied to batched operands.

03 Where does the loop body live after lax.scan, and what is the carry in that picture?

As a nested jaxpr inside the single scan equation’s params. The carry is the explicit state threaded from one iteration to the next, part of the equation rather than a hidden closure.

assigned

Readings