JAX

The library itself: tracing to jaxprs, jit and its cache, autodiff and batching as rule tables, pytrees, control flow, sharding, and the lowering to StableHLO.

Explainers

• What JAX Hashes Before It Decides Not to Compile: cache_key.py, Line by Line — Every jitted call ends at one question: have I compiled this exact thing before? Answering it means turning a compiler module, a set of devices, a pile of options and the environment your process happens to be in into one string. This is the 399-line file that does it, read whole at a pinned commit. Eight named ingredients go into a single running SHA-256, in a frozen order, and roughly a third of the file exists only to make the bytes arriving at those updates identical on two machines doing the same work: clone before mutating, strip debug info, then strip the debug info hiding inside a base64 string that the strip pass walks straight past, replace an unstable pointer with a constant, renumber the devices so a multi-process GPU job shares one key. Then twenty-six flag names are deliberately dropped, because they change what lands on your disk rather than what lands in the binary. The signature exhibit is the entry list itself, ported in order, with real digests you perturb one ingredient at a time; a second exhibit ports the flag filter so you can watch a dump flag get sorted in and then thrown away.
• Partial Evaluation: Splitting One Traced Call Into Two Programs — jax.grad hands you two programs out of one Python function: a forward pass that runs on the arguments you passed, and a linear pass that runs later when a cotangent arrives. Nobody wrote the second one. It comes out of partial evaluation, one file in jax/_src that gives every value a single bit, known or unknown, and propagates it with one rule: an operation is known only if every one of its inputs is. Known operations are executed during the trace; the rest become equations in a second jaxpr, together with the values they need from the first. Those crossing values are residuals, and they are exactly the activations a backward pass holds in memory, which is why rematerialisation is a predicate handed to the same splitter rather than a separate feature. The page reads the rule in the source, computes the known/unknown split on a small function you edit, and answers the rest of what the file settles: where a jaxpr gets its variables, why a constant is sometimes inlined and sometimes carried, and why the sine of a number you typed is still an equation inside jit.
• linear_util.py, Line by Line — Every JAX transform is written as a function that wraps another function, and this 504-line file is the object all of them stack on. A WrappedFun holds the original Python callable, a tuple of the transformations to apply around it, and one write-once cell per transformation for the metadata a transformation discovers while running but cannot return through the ordinary call. Read top to bottom, the file explains why the arguments meet the last-applied transformation first while the results meet the first-applied one first, why the object is hashable on its transformation stack but never on its cells, and why memoising a traced call needs a weak dictionary keyed on the raw function with a second dictionary inside it. Nine regions, plus a stack you push real transformations onto and then call.
• Writing JAX Notebooks That Show Their Work — A cell that prints a correct array proves nothing about what JAX did to produce it. This reads the machinery at a pinned commit of jax-ml/jax and lays out what a notebook can put on the record. A jitted call is four stations, not one step: the wrapper, the traced form, the lowered module, the compiled executable, each a plain object you can hold in a variable and interrogate, and the page walks them as a scene whose artefact list at every stop is read from the real method sets. What make_jaxpr actually is, which is the trace station under a friendlier name. The two dialects a lowering will print, and the ValueError for a third. Why the cost analysis before compiling says outright that it is looking at unoptimised code, and which one to read instead. The three stage timers JAX already publishes on a public event bus, with their verbatim log strings, so a nine-second cell becomes a trace time, a lowering time and a compile time. Then the case worth committing to: the print inside your function that fires once and goes quiet, walked call by call against the real cache-key ingredients, including the redefinition a re-run cell causes. The 125 configuration options split into the 115 you can scope and the ten you get one shot at, the three that raise if a backend is already up, where the IR dumps land and what the filenames mean, why a cache directory you configured correctly stays empty, and the five ordered categories JAX will use to explain its own cache miss, including the message it keeps for a function being redefined on the same line. Closing with the eight-cell order that survives a runtime restart, and what breaks when each cell is missing.
• Dispatch and Compilation: The Six Caches Under a jit Call — The first call to a jitted function takes a second. The second takes microseconds. Nothing about your code changed, so the difference is entirely bookkeeping: six caches stacked one under the other, each one catching a different kind of repetition, each with its own key. This page walks the stack from the C++ fast path down to the executable on disk, naming every cache by the variable that holds it and every key by the fields it actually compares, at one pinned commit of jax-ml/jax. Then it puts the stack under your hands: issue calls, change one thing at a time, and watch which cache answers and which ones never get asked. By the end, a surprise recompile is a question with a procedure, not a mystery.
• How JAX Builds a Backward Pass Out of a Forward One: ad.py, Line by Line — Forward-mode differentiation is easy to believe in: carry a derivative alongside every value and push both through the program together. Reverse mode, the one that makes training possible, looks like a different algorithm entirely. This 1162-line file says it is not. It runs the ordinary forward pass with one change, splitting each operation into a primal half that executes now and a linear half recorded as a second small program, and then it runs that second program backwards. Read whole at a pinned commit of jax-ml/jax, in the file order rather than the order things get called: the jvp entry point and its two-value tracer, the linearizer that builds the tangent jaxpr and prunes residuals it can forward instead of storing, the transpose interpreter that scans forward to find which equations touch a cotangent and then walks them in reverse, the three accumulator classes cotangents land in, the two registries a transpose rule can live in, the six helpers that install a derivative for a primitive, and the marker primitive whose transpose is your own custom_vjp function. The signature exhibit ports the real jvp and transpose rules for sin, mul and add and runs the whole pipeline on a function you pick: the primal equations, the tangent equations each rule emits, the linear jaxpr with its residuals named, then the reverse walk with every accumulation shown and the derivative it lands on checked against the analytic one. Closing on the counterfactual people get backwards: what actually happens when you differentiate a non-linear primitive that has no transpose rule at all.
• Where vmap Puts the Axis: batching.py, Line by Line — jax.vmap adds a dimension to your arrays and then spends the rest of the call deciding where that dimension should sit. This is the 789-line file that decides. Every value under a vmap is a pair, the array and one integer saying which of its axes is the batch, and every primitive gets a rule that answers one question: given where the batch axis is on the inputs, where is it on the output? Read the file and vmap stops being a vectoriser and becomes a table you can extend.
• How Every JAX Transform Unpacks Your Data: tree_util.py, Line by Line — jit, grad, vmap and shard_map all take whatever nested mess of dicts, lists and custom classes you hand them, and none of them contains a line of code for walking a container. They do not need one. Every transform calls the same pair of functions first: flatten the argument into a flat list of values plus a description of the shape it came out of, transform the flat list, put the shape back. This is the 1385-line file that defines that protocol, read at a pinned commit of jax-ml/jax, with the eleven regions that matter for flattening and registration walked line by line and the rest left in the pane. Four separate registries built at import, not one, and why a type registered through the public API lands in all four at once. What the C++ extension actually does when you call flatten, including the two rules people get wrong: the is_leaf predicate is consulted before the type lookup, on every node including the root, and dictionary children are visited in sorted key order, so the leaf order is not the order you wrote. Six ways to register a class and what each one costs. Why None is a node with zero children rather than a leaf, and what that does to a treedef. Why a treedef that prints identically to another can still compare unequal. The signature exhibit flattens a nested container you edit, computing the leaf list, the key paths and the treedef string the way the extension builds them.
• Control Flow as Primitives: One Trace, However Many Iterations — Write a Python for loop inside jit and JAX runs it at trace time, once per iteration, leaving ten thousand equations behind for a ten-thousand-step loop. The four functions in jax.lax that avoid that all work the same way: they trace the body exactly once, hang the resulting jaxpr on a single equation as a parameter, and hand XLA one loop op. The price is a rule you meet the first time you write one, and it is stricter than Python: the value carried from one iteration to the next must have the same pytree structure, the same shape and the same dtype every time round, with one narrow exception for Python scalars that costs a second trace. This page reads the real check and its error message out of the source at a pinned commit and hands you a carry-type checker you can edit, then follows the consequences: why both branches of a cond always run, why fori_loop is quietly two different primitives depending on how you passed the bounds, why scan can be differentiated in reverse and while_loop refuses outright, what happens to a branch under vmap, and what the unroll dial actually reshapes.
• shard_map Internals: The Per-Device Body and the Checks Around It — Everywhere else in JAX you write the whole-array program and let the partitioner work out the communication. shard_map inverts that: you write the program one device runs, on one shard, and you write the collectives yourself. This page reads the file that makes that work at one pinned commit of jax-ml/jax. What in_specs and out_specs each promise, the six checks that run before your function is ever called, the shape arithmetic that turns a global array into a per-device one and back, the type-level record of which mesh axes a value is allowed to differ along, which collectives are legal inside the body and what each one does to that record, and the single region the whole thing lowers to. The spec checker is on the page: pick a mesh, an in_specs, a body and an out_specs, and read the per-device shapes and the real verdict.
• From Jaxpr to StableHLO: One Rule Per Primitive — The last thing JAX does before handing your program to XLA is walk the jaxpr and, for every equation in it, look up a function in a dictionary keyed by the primitive. This page opens that dictionary: 311 primitives and 462 registrations counted from the call sites at one pinned commit of jax-ml/jax, filterable by platform, every row a link to the line that registers it. Then the wrapper around the loop, in the order the code builds it: the arguments you never passed and why they are in that order, the sharding and layout and donation attributes the partitioner actually reads, the tokens that give two debug prints an order, and the error you get when nobody wrote a rule for your accelerator.