jaxpr

A jaxpr is the recording tracing produced: a flat list of equations in single-assignment form. Each line binds one fresh variable to one primitive applied to earlier variables, with every shape and dtype stated. There is no control flow left, no module structure, no Python. It reads like assembly for tensor math, and it is the most honest description of your program you can get without leaving Python: print it with jax.make_jaxpr(fn)(*args).

Here is our attention function as a real jaxpr. Things Python never showed you are now explicit: k.T became a transpose equation, the matmuls became dot_general with their contraction dimensions spelled out, and the broadcast that aligned the row max for subtraction became a broadcast_in_dim with stated shapes.

{ lambda ; a:bf16[128,64] b:bf16[128,64] c:bf16[128,64]. let
    d:bf16[64,128] = transpose[permutation=(1, 0)] b
    e:bf16[128,128] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] a d
    f:bf16[128] = reduce_max[axes=(1,)] e
    g:bf16[128,1] = broadcast_in_dim[
      broadcast_dimensions=(0,)
      shape=(128, 1)
      sharding=None
    ] f
    h:bf16[128,128] = sub e g
    i:bf16[128,128] = exp h

Look at the convert_element_type pair around the reduce_sum in the second half: JAX upcasts the bf16 summation to f32 and converts back. Nobody wrote that. It is a numerics decision the framework made for you, and the jaxpr is the first place it becomes visible. Learning to read IRs is learning to notice decisions like this one, because one layer further down, you will be the one making them.

    j:f32[128,128] = convert_element_type[new_dtype=float32 weak_type=False] i
    k:f32[128] = reduce_sum[axes=(1,)] j
    l:f32[128,1] = broadcast_in_dim[
      broadcast_dimensions=(0,)
      shape=(128, 1)
      sharding=None
    ] k
    m:bf16[128,1] = convert_element_type[new_dtype=bfloat16 weak_type=False] l
    n:bf16[128,128] = div i m
    o:bf16[128,64] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] n c
  in (o,) }