the jax path · 0/12
start the path
the gym
jax floor

serves jax-path chapters 02 · 05 · streaks check themselves off; the same skill counts on every floor that drills it · all floors

Read the recording until it is boring.

The jax path's claim is that a jaxpr is readable on sight, and reading is a drilled skill. These stations run on the same generated corpus the kernel floor uses: seventeen real programs, traced and dumped, answers computed from the structured facts, never invented. Streaks persist in your browser and nowhere else.

drills

Active recall, one tap

GYM·02 the shape oracle
opdot_general
  • bfloat16[64, 128]
  • bfloat16[128, 512]
streak 0program: mlp

what shape and dtype comes out?
equations sampled from 278 real eqns · distractors computed, never invented · assigned in jax chapter 02
GYM·06 name the transform
the same function, traced under one transform. Which one?streak 0 / 5
{ 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
  in (g,) }
redacted blocks hide primitive names; shapes and structure are the tell
one function traced under six transforms · loud primitive names redacted for the game · assigned in jax chapter 05
the corpus
x-rayed

Seventeen real programs, wired together

Pick a program and hover any column: the source line, the jaxpr it traced to, and the StableHLO it lowered to light up as one. This is jax chapter 02's mastery bar made physical: narrate a program you have not seen, then another. When the middle column stops surprising you, the chapter has done its work.

GYM·04 the corpus x-ray
12 mapped equations · 26 of 30 jaxpr lines · 12 of 30 StableHLO lines

the transformer FFN: two matmuls around a GELU; watch bias broadcasts appear

source
jaxpr
StableHLO

bright lines are mapped equations · dim lines are scaffolding, compiler-made constants and broadcasts, or alignments the mapper skips rather than guesses

hover any line, bright or dim: it explains itself here · source lines recorded by jax itself, IR mapping computed mechanically
17 programs, 236 sync groups, all derived by the mapping script · nothing hand-drawn