Frameworks

How the array libraries actually work under the API: tracing a Python function into a typed intermediate form, the transforms (jit, grad, vmap) that rewrite it, the module systems that hold parameters, and the pytree contract that lets any of it compose.

Sub-topics

• PyTorch/XLA — The PyTorch frontend for XLA devices: lazy tensors, an HLO lowering, and the PJRT client that runs them on a TPU.
• 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

• Tracing → jaxpr: the One Trick Behind Every JAX Transform — jit, grad, and vmap look like magic, but they all do the same mundane thing first: they run your Python once, feeding it stand-in values that carry only a shape and a dtype, and record every operation into a small typed program called a jaxpr. Watch a real function become its jaxpr line by line, see the print that fires exactly once, learn why a new input shape makes jit retrace, why a plain Python if throws under jit, and how the pytree flatten/unflatten contract lets all of it compose. Every other JAX page rests on this one.
• jit(grad(vmap(f))): Why Transform Order Changes the Answer — jit, grad, and vmap all take a function and return a function, which makes them feel like decorators you can stack in any order. They are not: vmap(grad(f)) computes one gradient per example — the standard trick behind per-example gradients — while grad(vmap(f)) does not even type-check until you bolt on a reduction, and which reduction you pick changes the number you get. We ground it in JAX’s own tracer classes, reorder a live jit/grad/vmap chip stack and watch the trace-stack diagram and a real numeric example recompute together, then find the same order-sensitivity again in jacfwd vs jacrev.
• Params Are Data: the Functional Model Behind Flax — A Flax model never holds its own weights. Call init() and it hands back an ordinary nested dict of arrays; call apply() and it takes that same dict back in to run the forward pass. We follow one PRNG key as it splits down the module tree, watch every self.param drop a leaf into the growing params pytree, send that pytree back through apply, and catch the exact bug Flax’s own docs name when submodules get built inside a branch — the discipline that lets checkpointing, sharding, and jax.grad treat a whole model like any other array.
• The Seam Interface: computation_client.h, Line by Line — Everything PyTorch does on an XLA device leaves the process through one C++ class. Walk the real computation_client.h top to bottom: the four different things called "Computation", the transfer surface and its GIL warning, Compile, ExecuteReplicated, and the 43 pure-virtual methods that are the entire seam between PyTorch and PJRT.
• The Lazy Tensor: What Happens Between Your Op and sync() — On the xla device an op runs nothing. It appends a node to a graph, and the tensor you get back is a promise. This walks the whole mechanism in PyTorch/XLA at a pinned commit: what an XLATensor actually holds, how each IR node hashes itself as it is built and why shape only enters through the leaves, what torch_xla.sync() sweeps up and folds into one graph hash, the four ingredients of that hash in the order the executor folds them, the 2048-entry compile cache a hit or a miss lands in, and the reads (.item(), .cpu(), a stray print) that cut the graph somewhere you did not ask. Then the same machinery seen twice more: eager mode as a cut after every op, and torch.compile with the openxla backend as a cut pinned to the function boundary, traced once and replayed by hash.
• How PJRT_DEVICE Becomes a Client: pjrt_registry.cpp, Line by Line — Set PJRT_DEVICE=TPU, import torch_xla, and a chip you never named starts running your model. The whole decision is 174 lines of C++ in one file. We read all of it: the plugin interface with its three questions, the global map seeded with a single placeholder entry, the exact-string lookup whose error message doubles as the documentation, and then the if/else chain itself, branch by branch. The dynamic-plugin branch that a plain import turns on by default, and which dlopens a vendor .so, initializes it, and wraps it through the PJRT C API. The distributed key-value store that only exists for plugins that ask for it, built out of the same coordinator torch_xla uses for preemption. The CPU branch, the one device with no plugin at all. The TPU branch and its three-deep search for libtpu.so. The two half-finished branches, XPU and NEURON, that skip the plugin initialize and the profiler hook. And the else that catches your typo. The signature exhibit runs the real branch order: type a device string, flip the dynamic-plugin switch, and watch which branch catches it and which environment variables it reads.
• The PJRT Boundary: One Training Step, Crossing by Crossing — Somewhere between your PyTorch code and a TPU there is a line, and it is one C++ class: ComputationClient, pure virtual, forty-three methods, one live implementation. This page walks a single lazy training step across it four times (inputs down, program down, run, one value up), names the PJRT method waiting on the other side of each crossing, and counts them against a published lab capture on a Colab TPU. Then the things that never cross: an all-reduce is an instruction inside the program, not a call, and buffer donation rides down as an HLO annotation because the execute options struct is two booleans wide. Closes with the metrics report that has the seam running down its middle, the fourteen files that reach the client, the SPMD path that swaps one method for its plural, and the pinned OpenXLA commit that makes every line number checkable.
• Writing torch_xla Notebooks That Survive Colab — A Colab notebook that trains a model on a TPU is easy to get running once and surprisingly hard to get running twice. This is the craft layer under torch_xla: why torch and torch_xla are one pinned pair and never two independent versions; what PJRT_DEVICE actually does when you leave it unset (import picks a default and tells you in a warning line most people scroll past); why torchax and torch_xla cannot share a runtime, at run time over PyTorch’s dispatcher and again at install time over libtpu; why the first step is slow and the second one is not; and the three instruments that let you prove any of it: the metrics report, the IR and HLO dumps written to a file whose name is not the name you gave it, and the two calls that print a graph without running it. Closes with a cell order that works and a paste-back discipline that keeps a reference run honest.
• torch.compile Meets the Lazy Tensor: dynamo_bridge.py, Line by Line — Name openxla as your torch.compile backend and PyTorch hands the captured FX graph to one 794-line Python file. The obvious guess is that it lowers those nodes to XLA directly. It does not. It runs the graph once, on your real device tensors, through the same lazy runtime an un-compiled program uses, takes the hash of the recording, compiles under that hash without executing, restores every tensor the run modified, deletes the recording, and hands back a closure that holds the hash. We read all of it: the matcher that rebuilds the parameter list from your arguments and the trace-time weights, the three small classes that put duplicated, pass-through and None outputs back before the caller sees them, the tracing function and its five undo steps, the closure that runs on every call and never compiles, the collector that discovers unsupported operations by running the graph node by node and watching a counter, and the partitioner that quietly turns one compiled region into several XLA programs with host code between them. The signature exhibit runs the file’s own branch: press a call, watch the numeric cache key decide whether the trace runs, and watch the launch happen either way.
• Where torch_xla Calls PJRT: pjrt_computation_client.cpp, Line by Line — Everything PyTorch/XLA ever asks a device to do goes through one 1,073-line C++ file, and this walk reads all of it. Initialize, which runs once per process and quietly decides the names your devices will answer to. TransferToDevice, where a host tensor becomes a device buffer and an empty lambda is the only thing keeping the source alive. Compile, whose 143 lines are almost entirely two arms of an if: one that sets num_partitions to the device count and one that sets num_replicas to it, with a device assignment matrix transposed between them. ExecuteComputation and ExecuteReplicated, which take the same executable, disagree about strict shape checking, and lock in completely different ways. TransferFromDevice, the one call that genuinely blocks. And a first function in the file that is never called at all. The signature exhibit runs one step through the file in its own order, each crossing showing the PJRT method underneath it and the timer it stamps.
• SPMD in torch_xla: A Mesh, an Annotation, and One Virtual Device — Your model outgrew one chip. The usual PyTorch answer is to rewrite it into a parallel model; the SPMD answer is to keep writing single-device code and tell the compiler how a few tensors are cut across the chips. This walks the whole path in PyTorch/XLA at a pinned commit: what use_spmd() switches on and why every tensor then lands on a device called SPMD:0, what a Mesh and a partition spec really are, the permutation that turns a spec into a tile assignment, what mark_sharding does to the bytes you already uploaded, how the annotation reaches the HLO, and why one step becomes ExecuteReplicated behind a single device lock.
• From Pending IR to a Device Buffer: xla_graph_executor.cpp, Line by Line — On an XLA device your operations do not run when you write them. They pile up as unexecuted nodes, and then one call turns the pile into a single number. This is the 1,609-line C++ file that does it, read end to end. The compilation cache and the two environment variables that size it. The arena that knows every live tensor, and why a random seed has to be a small graph of its own. The six merge sites, scattered across nine hundred lines, that assemble the graph hash: a config flag a print sets differently from a sync, the two git revisions baked in at build time, one hash per synced tensor, the parameter order, the donated buffers, and the sharding mode. The lookup that decides compile or replay, and the two counters that tell you which happened. Buffer donation and the one condition under which it is safe, argued in the file’s own twenty-line counterexample. Then lowering, sharding annotations, parameter wrapping above 3200 inputs, and the single call that costs the seconds. The signature exhibit folds the real hash term by term, in the order the code folds it, so you can watch where one changed ingredient makes every later digest diverge.
• 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.