JAX Core, From the Inside
JAX has one trick: run your Python once with stand-ins for the arrays and record what it did. Everything else, jit, grad, vmap, sharding, is a program that reads or rewrites that record. The track starts with a jaxpr you print in a notebook and ends at the lowering that turns it into StableHLO, with the cache, the transforms, and the pytree protocol read in between.
Phase 1 — Run it first
Print a jaxpr, lower it, compile it, and read each artefact before any article names the code behind it.
- LAB·J1 · Trace, read, raw ↗ — Print jaxprs for a handful of functions and read what tracing kept and what it refused.
- Writing JAX Notebooks That Show Their Work — The four stations of a jitted call, the artefact each one hands you, and the cell order that survives a runtime restart.
- Tracing, the chapter ↗ — The kernels-site chapter that first shows shape-and-dtype tracing and what a tracer refuses.
Phase 2 — Tracing
The wrapped function, the tracer, and the partial evaluator that turns a call into a jaxpr.
- Tracing → jaxpr: the One Trick Behind Every JAX Transform — Run a Python function once with stand-ins and you get a typed jaxpr; the concept page that starts the JAX stack guide too.
- Partial Evaluation: Splitting One Traced Call Into Two Programs — One bit per value decides which half of your function runs now; residuals are what crosses, and they are the activations a backward pass keeps.
- linear_util.py, Line by Line — The object every transform is applied to: WrappedFun, the stack that wraps it, and the write-once cell each transformation gets.
- jax/_src/core.py ↗ — Tracers, primitives, jaxprs, and the type system, in one large file the pages above cite by line.
- The JAX machine ↗ — The kernels-site chapter on the tracer as a machine.
Phase 3 — jit and the cache
What a jit boundary is, what goes into the cache key, and where a recompile comes from.
- What JAX Hashes Before It Decides Not to Compile: cache_key.py, Line by Line — cache_key.py read whole: the eight ingredients hashed in order, the canonicalisation before them, and the twenty-six flags left out.
- Dispatch and Compilation: The Six Caches Under a jit Call — The six caches between a jitted call and a PJRT executable, each named by the variable that holds it and keyed on what it actually compares.
- LAB·J2 · The recompile hunt ↗ — A training loop that recompiles every step; find why from the counters, not from guessing.
- jax/_src/dispatch.py ↗ — The dispatch path a call takes after the cache decides.
- jax/_src/pjit.py ↗ — jit itself: the staging, the sharding resolution, the caches.
- jax/_src/compilation_cache.py ↗ — The persistent cache on disk and what it stores.
- The cache, the chapter ↗ — The kernels-site chapter on the two counters and the price of static.
Phase 4 — Transforms
grad and vmap are tables of rules, one per primitive; read the tables and the transforms stop being magic.
- jit(grad(vmap(f))): Why Transform Order Changes the Answer — Why vmap(grad(f)) and grad(vmap(f)) compute different things, and the order that means what you think.
- How JAX Builds a Backward Pass Out of a Forward One: ad.py, Line by Line — ad.py read whole: the two-value tracer, the linearizer that emits a tangent jaxpr, and the pass that runs that jaxpr backwards to produce gradients.
- Where vmap Puts the Axis: batching.py, Line by Line — batching.py read whole: the tracer that carries one integer per value, the rule tables that move it, and the two branches behind the missing-rule error.
- LAB·J3 · Own the derivative ↗ — A custom_vjp written from the rule, checked against autodiff, then a remat trade-off measured.
- jax/_src/custom_derivatives.py ↗ — custom_jvp and custom_vjp: how a user rule enters the tables.
- jax/_src/ad_checkpoint.py ↗ — jax.checkpoint: the remat policies and what they keep.
- vmap, the chapter ↗ — The kernels-site chapter on where the batch axis goes.
Phase 5 — Structure
Pytrees are how arbitrary Python data crosses a transform; control flow is a handful of primitives with their own rules.
- How Every JAX Transform Unpacks Your Data: tree_util.py, Line by Line — Every transform flattens your container first; this is the file that defines what a leaf is, what a treedef remembers, and the six ways to register a type.
- Control Flow as Primitives: One Trace, However Many Iterations — One trace however many iterations, the carry as a type checked by a real function, and why while_loop refuses a gradient.
- jax/_src/lax/control_flow/loops.py ↗ — scan and while_loop, with their batching and differentiation rules.
- Pytrees, the chapter ↗ — The kernels-site chapter on flattening and treedefs.
- Control flow, the chapter ↗ — The kernels-site chapter on both branches and one index.
Phase 6 — Devices
A mesh, a sharding, shard_map, and the lowering that produces the StableHLO XLA partitions.
- Sharding in JAX — Name your devices into a mesh, annotate an array, and the compiler places the bytes; the design-level page.
- shard_map Internals: The Per-Device Body and the Checks Around It — The manual mode from the inside: the specs it checks, the per-device shapes it computes, and the one region it lowers to.
- From Jaxpr to StableHLO: One Rule Per Primitive — The rule table read whole: 311 primitives, per-platform variants, the equation cache, and the attributes the partitioner reads off the entry function.
- LAB·J4 · Eight fake devices ↗ — Shard an array three ways on simulated devices and read the collectives the partitioner inserted.
- jax/_src/xla_bridge.py ↗ — The PJRT clients JAX constructs and the plugin discovery.
- jax/_src/mesh.py ↗ — Mesh and its axis names.
- Sharding, the chapter ↗ — The kernels-site chapter on mesh, spec, and placing an array.