Tracing turns that block into this, with one call to jax.nn.softmax already spread across several primitives.
j:f32[] = sqrt 32.0
k:f32[16,16] = div i j
l:f32[16] = reduce_max[axes=(1,)] k
m:f32[16] = max -inf l
n:f32[16,1] = broadcast_in_dim[
broadcast_dimensions=(0,)
shape=(16, 1)
sharding=None
] m
o:f32[16,1] = stop_gradient nReading a 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 hThe detail worth staring at
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,) }Lessons
- 01The anatomy of a jaxprA ClosedJaxpr is a jaxpr plus the values it closed over, and the grammar underneath is small enough to hold whole. ·
- 02Transforms are jaxpr rewritesgrad, vmap, scan, remat, shard_map: one small dense function, traced once, then rewritten five ways in front of you. ·
- 03Reading nested jaxprsSome equations carry whole programs in their params. The indented block underneath is a jaxpr like any other, and kernels nest exactly this way. ·
- 04Params, and the drills that fix themParams are the configuration a primitive needs before it can run. Fluency is repetition, and the gym exists for exactly that. ·
- 05One name, one definitionPython let you assign to x three times. The jaxpr that came back has four names, and the rule behind that is the one property every IR below this one also keeps. ·