a mastery course · the jax path · across the language
Learn to think in JAX.
JAX is one idea worn twelve ways: your function becomes data, and everything else (grad, vmap, jit, sharding) is a transformation of the recording. Part i builds that model until surprises stop. Part ii runs it at production scale.
your progress 0/12 chapters
How JAX thinks. Eight chapters that make the machine's behavior derivable instead of memorized.
- ch 01 Arrays An array in JAX is a value you compute with, not memory you overwrite.
- ch 02 Tracing Once jit calls your function, it stops running Python and starts recording a program.
- ch 03 Jit A jitted function is compiled once per signature, then reused exactly until that signature changes.
- ch 04 Autodiff `jax.grad` does not estimate a slope; it derives one exactly from the recorded program, and which direction it walks that program is a cost decision you get to make.
- ch 05 Vmap `vmap` doesn't loop over your data; it rewrites the program itself to already expect a batch.
- ch 06 Control flow The trace can only see control flow you hand it in a form it understands, and everything else either breaks or silently unrolls.
- ch 07 Pytrees A pytree is not a data structure you chose. It is the contract every transformation shares.
- ch 08 Randomness JAX has no global random state, only keys you pass by hand.
State, the mesh, honest measurement, and the training run that composes all of it.
- ch 09 State In JAX, state is not something a class holds. It is something a function returns.
- ch 10 Sharding A sharded array is still one array; the mesh and the spec just say where its pieces live.
- ch 11 Performance Most JAX benchmarks measure the wrong thing without anyone noticing.
- ch 12 The Training Run Every fact this track taught shows up once more here, wired together into a loop that actually trains.
Every page of the official JAX tutorial series is assigned reading somewhere on this path, and the chapters keep going where the docs stop: the cache key, the cost model of reverse mode, the sharding ladder, the benchmarking method.
Every snippet ran before it was printed (jax 0.4.38 on CPU, this machine, 2026-07-27; eight host devices for the sharding chapter), and printed outputs are real outputs. Where the current API differs from the verified one, the caption says so.
When a chapter names XLA, fusion, or the interconnect, the kernel path teaches that layer with its own hands: fifteen chapters down the stack, gates measured on real TPUs. The two paths cross-reference; neither repeats the other.