AI / ML

How machine-learning systems are built and made fast — training, inference, and the plumbing that moves tensors and weights around.

Sub-topics

• Fundamentals — The architecture under every modern model, built from first principles: tokens and embeddings, attention computed number by number, multi-head and the KV cache, positional encoding, and the block stacked into a full network.
• Training at Scale — Spreading one model across thousands of GPUs: data, tensor, and pipeline parallelism, FSDP/ZeRO sharding, activation checkpointing, and the interconnect that decides which of them you can afford.
• Checkpointing — Saving and loading model state at scale: sharded arrays, multihost clusters, and reading exactly the bytes each machine needs.
• Post-Training — Turning a base model into an assistant after pretraining: reward signals, group-relative advantages, and the RL clusters — rollout, reference, reward, and the learner — that keep a policy, its sampler, and its judges in sync.
• Inference & Serving — Turning a trained model into a service that answers thousands of streams at once: the memory-bandwidth roofline, the KV cache, continuous batching, paged attention, and the scheduler decisions that decide latency and cost per GPU-hour.
• Applied Systems — The production plumbing that puts models in front of traffic: feature stores, serving paths, and the data systems that keep training and inference reading the same values under real load.
• Compilers & Kernels — What happens between your model code and the silicon: the compiler that splits one program across a device mesh, decides which operations fuse, assigns every buffer an address, and emits the collectives that move tensors between chips.
• 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.
• Agent Infrastructure — The plumbing that lets a model act: tool catalogs and the router that picks from them, the credentials an agent borrows on your behalf, the events that wake it up, and the sandboxes where its code runs.

Explainers

• Design an Embedding Retrieval System — Search where the query and the item share no words at all — a photo finds the listing, a phrase finds the video, a person finds the people they might know. The retrieval shape under visual search, video search, similar listings, and People-You-May-Know, built from zero: what an embedding is (a learned point in space; nearness means similarity), why two towers instead of one model, a commit-first envelope on a 100M-item index, the honest deliberation between a brute-force GPU scan, IVF partitions, and an HNSW graph — with recall vs latency computed live on a real toy corpus — then the offline indexing pipeline, the online path, keeping the index fresh, training the towers on engagement with in-batch negatives, the multimodal shared space, and the failure sweep whose quietest box returns a confident 200 while silently retrieving nonsense.
• Evals & Experimental Design — Most ML claims die at one question: compared to what? This is the research-taste page — the difference between "my method gets 71.2%" and "my method is actually better." We walk the claim-evidence gap, the tuned-baseline rule, the eval hierarchy from loss to production, and benchmark contamination — then the part everyone skips: is 71.2% really better than 70.8%? Dial an interactive significance explorer with real binomial confidence intervals and a live bootstrap, watch the variance a single model has run-to-run, judge the LLM judges honestly, and design the experiment before you run it.
• Optimization Dynamics: Why Adam, Why Warmup, Why Cosine — An optimizer never sees the loss surface — it gets one gradient at a time and has to turn that stream of local hints into a path to the bottom. We race SGD, momentum, and Adam down the same dialable ill-conditioned valley (real update rules, trajectories drawn live), then build Adam precisely — per-parameter learning rates from the second moment, bias correction honestly — and AdamW's decoupled decay. Then the practitioner's layer: why warmup exists, what cosine/linear/WSD actually buy, gradient clipping and loss spikes, muP transfer, and batch-size-vs-LR — all computed, all animated.
• Mixture of Experts, Routed Honestly — A dense model spends every parameter on every token; a mixture-of-experts model keeps far more parameters but touches only a slice per token — huge in memory, cheap in compute at once. The whole design turns on one question: which experts does a token go to? We compute the routing decision live — a real softmax gate over 8 experts, top-k selection, a capacity cap that drops overflow tokens, and the load-balancing loss that fights router collapse — then count Mixtral 8x7B down from 47B parameters to the 13B it actually uses, cover expert parallelism and its all-to-all bill, and the DeepSeek fine-grained-plus-shared recipe.
• Design a Recommendation System (in the LLM Era) — Millions of items, one person, a hundred-millisecond budget — you can never score the whole catalog, so every large recommender is a funnel: cheap retrieval narrows millions to hundreds, a heavy ranker scores only the survivors, and a re-ranker shapes the final slate. We compute the scoring budget, dial an interactive funnel until it blows past 100ms, walk two-tower retrieval and the feature store behind the ranker, face the feedback loop that quietly collapses the catalog, and separate what the LLM era really changes from the hype — grounded in YouTube and Instagram's own papers.
• How to Design an ML System in 45 Minutes — An ML design prompt sounds like "build a model," but the room is grading something narrower: can you turn a fuzzy product goal into a crisp ML objective? That one translation is load-bearing — "reduce harmful content" can become a classifier with a review queue, a risk ranker for a bounded queue, or a rationale generator, and each spawns a different system. We make that fork the signature exhibit, then walk the rest of the round the way an interviewer scores it: where labels really come from, why the first model should be simple, why offline wins so often lose online, and the retraining loop that keeps a live model honest — grounded in Google's Rules of ML.
• Design a Ranking System — Ten million items, one person, two hundred milliseconds — you can never score the whole catalog with your best model, so every large ranker is a funnel. We compute the scoring budget first (a 5 ms/item heavy model against a 200 ms wall scores ~40 items, not ten million), then build the funnel live as a budget negotiation, walk candidate generation → light ranker → heavy ranker → re-ranking, and face the parts interviews skip: calibration (why an uncalibrated ad model loses real money), position bias in the training logs, offline-vs-online divergence, and the harmful-content moderation variant with its human-review queue — grounded in the published ad-CTR papers.
• ML Reliability in Production — When a web service breaks it throws a 500 and pages you; when a model breaks it keeps answering — confidently, wrongly, and silently. We name the four ways models rot (data drift, concept drift, upstream schema breaks, feature-pipeline rot), build a live drift simulator that computes PSI and KL off two real distributions and watches them trip their alarms at different speeds, then install the whole defense: monitoring the inputs and the outputs, evals on a schedule, the retrain-validate-shadow-promote loop, rollback when the model is stateful, and an incident runbook that triages bad model vs bad data vs bad code — drawn, computed, and animated.
• GRPO Advantage: Z-Score Your Siblings, Line by Line — PPO learns a whole second neural network just to guess how good an answer is. GRPO throws that away: sample a group of answers to the same prompt, and each answer's advantage is just how far above or below its siblings it scored. We walk the real tunix advantage estimators — GRPO's z-score, Dr.GRPO's un-normalized fix, and RLOO's leave-one-out baseline — three self-contained functions, computed live side by side.
• GPU Arithmetic: the Numbers That Decide ML Systems — Three numbers decide every ML-infra design: FLOPs to compute, bytes to move, bytes to store — everything else derives. We read the A100/H100/H200 spec sheet honestly (the sparsity asterisk quietly doubles the headline), plot dialable workloads on a real machine's roofline, derive the matmul 2mnk and training 6ND rules once, reproduce Llama-3's GPU-hours, and end on a live estimator: params, tokens, GPUs, MFU, dtype → days and dollars. Drawn, computed, and animated.
• What "Atomic" Means on a Filesystem vs. an Object Store — A checkpoint must appear all-or-nothing, even if the job dies with one byte left to write. We walk the real Orbax code that guarantees it two ways — an atomic rename on a POSIX filesystem, and a commit_success.txt marker on an object store that has no rename at all.
• Zero-RPC Sharding: How 1,000 Hosts Agree Who Writes Which Bytes — At checkpoint time every data-parallel replica holds the same shard — so who saves it? We walk the real Orbax planner that assigns ownership with zero coordination: each host runs one deterministic function over the sharding it already has, and every unique shard gets exactly one writer. Single-replica, replica-parallel, and the divisibility check that decides between them.
• The B-Tree Merge: Thousands of Checkpoint Files into One Atomic Manifest — Every host in a training run writes its own OCDBT checkpoint files into a per-process subdirectory — no coordination, no contention. Turning those scattered per-process B-trees into one global manifest a reader can trust is a single atomic TensorStore transaction. We walk the real Orbax file that does it: glob the subdirs, stage every copy_range under ts.Transaction(atomic=True), validate the would-be-merged store, then commit once — all of it, or none of it.
• Quantization for Deployment — A trained model is a pile of fp16 numbers, and at inference the whole game is how many bytes you drag from memory per token — because decode is bandwidth-bound. Quantization shrinks each weight to int8 or int4: a scale, a zero-point, a rounding, and the error that leaks. We compute it live on a real weight row, show why activations break at int8, walk GPTQ and AWQ honestly, dial the size-vs-accuracy trade, and cover KV-cache quant, FP8/NVFP4, and ternary BitNet — drawn, computed, and animated.
• How Qwix Quantizes Any Flax Model Without Touching Its Code — You have a trained Flax model and you want it in int8 — but the model code isn't yours to edit. Qwix quantizes it anyway, by a trick that has nothing to do with the model: it swaps the ops the model calls. For the duration of one forward pass it patches jax.lax.dot_general and jnp.einsum, so every matmul routes through a quantized version while the model's own source never changes. We drive the op-swap live, compute int8/int4/fp8/nf4 grids from Qwix's real bounds, show how PTQ and QAT ride the exact same mechanism, and map what gets quantized where — plus the honest costs of patching a running program.
• GPTQ, Line by Line: Hessian-Based Weight Quantization — Round a weight to fewer bits and you commit an error — GPTQ's move is to make the weights you have not rounded yet absorb it. We walk Google's real JAX implementation (qwix's gptq_core.py, 220 lines) top to bottom: why the Hessian is just X·Xᵀ, the dampening that keeps a Cholesky alive, the three-line factorization dance that replaces a matrix inverse, and the two nested loops where each column's error becomes a rank-one update on everything to its right. Then the algorithm runs live — scrub the blocksize and damping and watch the compensation wave move.
• The Post-Training Pipeline (SFT → RLHF → DPO) — A base model only predicts the next token — it has never been told to be helpful. Post-training is how it learns: supervised fine-tuning on demonstrations, a reward model distilled from pairwise preferences, then RLHF where four models fight in memory on a KL leash — or DPO, which folds the reward model into a single classification loss. We draw the reward curve, animate the four-model dance per batch, watch a policy hack a length-biased reward, and lay SFT/DPO/PPO/GRPO flat on a trade-off table.
• Scaling Laws, Honestly — Language-model loss falls as a clean power law in compute — which turns "how big a model?" from taste into arithmetic. Follow Kaplan's parameter race into Chinchilla's correction (tokens scale WITH parameters, about 20 per one), dial a FLOPs budget to its compute-optimal split, plot GPT-3, Chinchilla and Llama 3 honestly on the line, and see why inference economics now makes everyone over-train — with an honest beat on whether "emergent abilities" are real.
• Diffusion Models, Honestly — You can't paint a photo in one stroke, but you can clean up a slightly-noisy one — and if you can do that, you can generate. That is the whole trick. We destroy a real 2-D distribution on the DDPM noise schedule (the ᾱ_t math, computed live as you scrub time), state the ε-prediction objective in one honest equation, count the network evaluations a 50-step image really costs, work the arithmetic that makes latent diffusion 48× cheaper than pixels, dial classifier-free guidance between fidelity and diversity, and end on the honest history — why the field left GANs behind — and the failure modes nobody screenshots.
• Design a RAG System — A model with a knowledge cutoff and no access to your private documents has to answer questions about both — so you retrieve the right passages and hand them to it as context. We size the corpus in bytes and vectors, dial an interactive chunker until retrieval hits, draw the vector index and reranker honestly, and end on the part everyone skips: how you measure whether any of it worked. Drawn, computed, and animated.
• 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.
• Sharding in JAX — A 70-billion-parameter model does not fit on one accelerator, so JAX spreads each array across many — without you rewriting the math. Name your devices into a mesh, say which array dimension rides which mesh axis with a PartitionSpec, and the compiler places the bytes and inserts the communication. We build it up from one device to a full mesh, watch what jit does with a sharding — propagation, the gaps the compiler fills, the collectives it drops in — weigh shard_map against auto-sharding honestly, and face the Shardy cutover. Drawn, computed, and animated.
• The SPMD Partitioner: One Program Across a Device Mesh — You write a program as if one giant machine ran it, annotate a few tensors with how they are split across a mesh of chips, and the compiler rewrites it into N identical per-device programs — inserting every all-gather, reduce-scatter, all-to-all, and halo exchange needed to keep the maths correct. We trace how those annotations propagate through the graph, where each collective is born (and which resharding forces it), the one windowed op that needs halo exchange, and the case the partitioner cannot solve on its own — drawn, computed, and animated.
• The Checkpoint Lifecycle: What Happens Between save() and Durable — A training run that saves synchronously leaves a rack of accelerators idle for the seconds it takes to push a terabyte to storage. So the save goes async — a fast device-to-host copy, then a background thread that writes, coordinates every host, commits the directory as one atomic unit, and garbage-collects the old ones. Follow one save from the training step that triggers it to the moment it is durable and restorable onto a different mesh — drawn, computed, and animated.
• Anatomy of a FlashAttention Kernel — Plain attention is not slow because it does too much math — it is slow because it writes an N×N score matrix out to memory and reads it back, twice. FlashAttention never writes that matrix at all: it walks the keys in tiles, keeps a running max and denominator so softmax stays exact block by block, and recomputes what it needs in the backward pass instead of storing it. We compute the materialization bill, step through the online-softmax recurrence on a real tiny attention, do the IO complexity honestly, and cover what FA2 and FA3 actually changed — drawn, computed, and animated.
• Inside XLA: ~200 Passes and the Fusion Decision — A traced program enters XLA as a graph of abstract operations and leaves as a handful of GPU kernels. In between, a couple hundred passes rewrite it — and one of them, fusion, decides which operations share a kernel by putting a stopwatch on every merge. We walk the ordered pass pipeline, the memory round-trip fusion deletes, the priority cost model (time_unfused − time_fused) and its greedy queue, the vetoes that refuse a merge, buffer assignment, codegen, and the recompilation you pay for all of it — drawn, computed, and animated.
• Design an LLM Serving Platform — One trained model, thousands of concurrent chat and batch streams, and a GPU that costs by the hour. Start from the roofline that rules everything — prefill is compute-bound, decode is memory-bandwidth-bound — then build up through the KV cache, continuous batching, paged attention, and prefill/decode disaggregation. Drawn, computed, and animated.
• Design a Text-to-Image Service — The image-generation sibling of the LLM serving platform — and it inverts almost everything. A chatbot streams tokens the reader consumes as they arrive; an image is all-or-nothing for whole seconds, so the product is a queue with a progress bar, not a stream. Build it up honestly: what a request really is (prompt → text encoder → N denoising steps → decode, with the diffusion loop as a black box), a commit-first estimate of seconds-and-cents per image, the three serving shapes deliberated interviewer-style, the GPU fleet alive under real queueing arithmetic, LoRA adapters for the personalized-headshot product, two-sided safety with an honest fail-open/fail-closed call, the failure sweep, and text-to-video as the same skeleton at brutal scale.
• The Transformer, End to End — The researcher-round staple: draw the architecture from memory and explain every tensor. We build it once, honestly — token → embedding → attention computed on four real tokens (Q/K/V, a QKᵀ heatmap, softmax, the weighted sum, every number real) → multi-head → the KV cache and the MHA/MQA/GQA/MLA memory ladder dialed live → RoPE as a rotation → the MLP and residual highway → norm placement → the block stacked N times → a parameter count that reproduces Llama-3-8B to the last billion.
• Distributed Training, End to End — A 70-billion-parameter model needs 1.1 TB just to hold its weights, gradients, and Adam states — fourteen times what fits on one GPU. Compute the bytes, then install every fix in order: data parallelism with ring all-reduce, FSDP/ZeRO sharding, tensor and pipeline parallelism with their bubbles, activation checkpointing, and the interconnect that decides which of them you can afford — drawn, computed, and animated.
• Design a Feature Store — A model needs the same number at 3am training time and at the 50-millisecond serving moment — the user's 7-day click count had better mean the same thing in both. Build up from every team hand-rolling features in their service to a dual store fed by one definition: a columnar offline store holding years of history for point-in-time-correct training, a key-value online store holding the latest vector for 50K QPS under 50ms p99. Train/serve skew, the time-travel join, streaming freshness, and drift — drawn, computed, and animated.
• Loading a Safetensors Checkpoint on a Multihost Cluster, Line by Line — A safetensors file is one flat byte blob per tensor, but a big model must land as sharded jax.Arrays across many hosts. Walk the real loader that maps each host’s shards to byte ranges — and reads exactly those.
• Model Surgery: Rewriting a Checkpoint’s Parameters, Line by Line — A model is saved under one structure and needed under another — HuggingFace names, a split gate/up projection, 64 separately-saved experts. We walk Orbax’s real model-surgery module: six transforms that rename, fuse, repeat, and stack a pytree of weights, and the host-memory care that lets them run on tens of gigabytes.
• The RL Cluster: Five Roles, One Mesh Dial — An RL policy under training needs five jobs sharing one accelerator fleet: sample completions, score them, hold a frozen baseline, and push a gradient update back — and one line of config decides whether those jobs share memory or ship weights across the wire. We open Google's real tunix RL cluster to see exactly how, and compute what a weight sync costs either way.
• Anatomy of a Production JAX LLM Trainer — One training loop, five subsystems that all have to agree: a device mesh, a table of sharding rules, a jitted train step, an asynchronous checkpointer, and a data pipeline that never hands two devices the same example. We build the mesh from four integers, watch a logical axis resolve onto physical hardware, dial a live mesh and shard a toy tensor, then follow one step through jit, one checkpoint through an async write, and one host through a crash — the whole trainer, end to end.
• 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.
• Design a Tool-Calling Platform for AI Agents — Thousands of tenants, each running agents that need to act on a real SaaS account on a real end user's behalf — a shared connector catalog, a token vault, and an execution plane that survives one provider having a bad day.
• The Tool Router: the Right 5 Tools out of 10,000 — An agent platform can wire up thousands of tools; one model call has one context window. We measure what a tool actually costs in tokens, price what ten thousand of them costs per call, and walk five ways platforms narrow the catalog down to the handful that matter — plus how to tell whether the narrowing actually worked.
• 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.
• The Loop That Runs Every XLA Pass: hlo_pass_pipeline.cc, Line by Line — A compiler pass is small and well behaved: it takes a program, makes one kind of improvement, and says whether it found anything. Running a few hundred of them back to back is a different job, and this is the 334-line C++ file that does it, read end to end. The two flags that skip passes, their four-shape grammar, and the two occurrence counters they resolve against. The invariant checkers between passes, why they only run after a pass that changed something, and the single rule that makes them safe. The debug mode that hashes the whole module twice to catch a pass that lied about its own return value. The dump condition with a special case hiding inside it, where the catch-all pattern writes fewer files than a narrow one. And the written ledger every pass appends a row to, opened before the skip decision so its numbering survives a bisect. The signature exhibit is the loop itself, ported line for line over a seven-pass pipeline, so you can watch a flag reshape a run.
• The PJRT Plugin Contract: pjrt_api.cc, Line by Line — A vendor ships a shared library for an accelerator the framework has never been compiled against, and it runs your model anyway. The entire arrangement is one exported C symbol and two integers, and this is the 209-line file where the two sides meet. We read all of it: the global map from a device-type string to a function table, the lowercase that is the whole matching rule, the write-once registration that refuses to let two packages claim one name, and the loader itself, which dlopens a library, asks for a single symbol called GetPjrtApi, and never closes the file again because the table lives inside it. Then the handshake. Where the plugin gets its version number from, why major must match exactly and minor only has a floor, the environment variable that switches the rule from generous to strict, and the value that reads as encouraging while doing the opposite. The signature exhibit runs the real comparison order: set the version the plugin was built against, set the environment variable, and read back the exact error string absl::StrCat assembles, trailing-parenthesis quirk included.
• HLO Module Anatomy: What XLA Holds While It Compiles — An HLO dump is an object graph printed. This walks that graph in the XLA source at a pinned commit: an HloModule that owns a vector of computations and points at one of them as the entry, an HloComputation whose parameters and return value are both just instructions in its own body, and an HloInstruction that owns nothing beyond an opcode, a shape and a list of pointers out to operands and back from users. Then the two things that sit beside the graph and decide most of what it costs: the shape that carries a layout, where minor_to_major turns an index into an address and a disagreement turns into a copy you never wrote; and the schedule, an optional total order added late to hold peak memory down. It closes on the verifier that runs between all two hundred passes, in the order it runs its checks, and the refusals worth knowing by name.
• Reading an XLA Dump: What the Compiler Writes, and How to Read It — Someone asks you to attach the HLO dump. You set one flag, point a directory at it, and get back a pile of files with names like module_0000.jit_step.0007.simplification.after_algsimp.before_reshape-mover.txt. This reads that directory end to end at a pinned commit of openxla/xla: the four inference rules one flag triggers before anything is written, the filename grammar and what every segment of it means, the two named bookends of a compile and the four buffer-assignment reports that ride with the second one, the per-pass files and the rule that quietly drops any pass that changed nothing, and the flags that make the dump larger, smaller, or unreadable. Then the two tools that take a dumped file and do something with it without a framework anywhere in sight: hlo-opt, which re-runs named passes and prints any stage of the compile, and run_hlo_module, which executes the module and checks it against the interpreter. The signature exhibit assembles the actual directory listing from the real inference rules, one flag at a time.
• How a StableHLO Module Becomes HLO: mlir_to_hlo.cc, Line by Line — Your framework does not hand the compiler a graph. It hands it a versioned document written in a dialect called StableHLO, and 434 lines of C++ decide what that document is allowed to say. We read all of them. The includes, which split into the two vocabularies this file sits between. The thirteen MLIR passes a module walks through before anything called HLO exists: seven from a Shardy pipeline it pulls in unconditionally, six added here by name, ending with the one that copies captured constants into control-flow regions because XLA has no implicit capture. The Shardy-loses-to-GSPMD fallback at the top of the function, and the block sixty lines later that silently reads the flag it just switched off. The parser, and the one error message that names three different version problems because it genuinely cannot tell them apart. Then the write side, which is the longer half: the version clamp that keeps the lower of what a plugin asked for and what this build can produce, the allow-list that refuses to serialize any operation nobody has promised to read, and the twelve-week compatibility window written as a requirement rather than a number. The signature exhibit runs the real branches: pick a direction, flip the flags, and every pass your module walks through appears in order, at the line that adds it.
• Layout Assignment: Where the Copy in Your HLO Dump Comes From — A copy instruction shows up in your dump between two lines you wrote next to each other, and the only thing that changed across it is the little list in braces. This reads the pass that put it there, at a pinned commit of openxla/xla. What a layout is, which is one permutation of dimension numbers read from the fastest-moving end. The three kinds of constraint the pass collects, and why exactly one of them can be settled by inserting a copy. The switch statement that answers, for all 134 opcodes, whether an operation may hold one arrangement out and a different one in, and the 89 that may not, which is what lets a constraint travel across most of a graph for free. The deque with two ends, the three propagators hanging off it, and the cap of two rewrites that stops it oscillating. The derivation at a transpose, whose stated goal is to make the transpose move no bytes at all, and the copy that lands one edge above it as a result. Then the priority arithmetic that settles a collision, the row-major default handed to whatever nobody argued about, the tuple rule that sends a constraint further upstream instead, and the moment inside AssignLayouts where the copy is finally created. The signature exhibit runs the real propagation loop over a four-instruction module, one pop at a time, with every layout on screen produced by the same arithmetic the pass uses.
• Buffer Assignment: How Every Value Gets an Address — A compiled XLA program never allocates. It receives a handful of slabs from the runtime and every kernel was compiled with its offsets inside those slabs already baked in. This walks the pass that picks them, at a pinned commit of openxla/xla: the values that must share memory because the language says so, the schedule that turns a dependency graph into a clock so that a live range can be two integers, the closed-interval interference test and the one endpoint case it makes an exception for, the colour that is a memory space rather than a graph colouring, the four kinds of value that never reach the packing, and the biggest-first best-fit heap that packs everything left. Then the nine ordered reasons an allocation refuses a buffer, the in-place update whose copy you never see until something reads the old value, and the four files a dump writes with the answer in them. The signature exhibit packs a real schedule with the real rules and lets you drag a live range until the sharing breaks.
• What Shares a GPU Kernel: priority_fusion.cc, Line by Line — A GPU kernel launch costs about a microsecond of nothing happening, plus a round trip through memory for its inputs and another for its outputs. So the largest single decision an ML compiler makes is which operations get to share a kernel, and on the GPU backend that decision is 1,419 lines of C++ we read end to end. The type choice the whole pass turns on: a priority is an absl::Duration, the wall-clock time this merge is estimated to save, which is what lets an ordered map double as a priority queue and two infinities put bitcasts first and constants last with no special case in the ordering. The correction to the usual summary: the queue holds producers, not edges, and a producer is scored against every one of its consumers at once, so the cost of being duplicated into three kernels is always in the number. The incremental machinery that keeps the pass from being quadratic, where a re-score is the old score plus the delta from the new consumers minus the runtimes of the departed ones. The Triton path that is tried first and the elemental checks that run when it declines. The thirteen refusal strings, quoted verbatim, that are what you actually read in a fusion dump. And the three flags that let you watch it decide on your own model, including the compiler fuel that bisects to the exact merge that changed your numbers. The signature exhibit runs the real ordering rules over a small graph: press pop, watch a producer that scored below zero turn profitable because a neighbour was absorbed somewhere else.
• Instruction Fusion Legality: instruction_fusion.cc, Line by Line — Merging two operations into one kernel deletes a write and a read of a whole tensor, which on this hardware is the difference that matters. So the compiler wants to fuse everything, and this 1,248-line C++ file is where it works out what it may. The unit of decision is one graph edge: a producer, a consumer that takes it as an operand, and one boolean. Read end to end: the hand-written classification of all 134 opcodes into cheap and expensive, in a switch with no default case; the two opcodes that may always be copied and why copying them lowers memory traffic; the global pre-pass that bans a producer from duplication unless every consumer will swallow it; the reverse-post-order queue and the thirty-line comment about the duplicate clone it exists to prevent; and the quarter of the file that is about correctness rather than speed, where a slice read out of a buffer meets an update written back into it. The signature exhibit is the six-gate decision ladder ported from the file, run over one small module, so you can watch an edge get refused and see which predicate did it.
• IFRT Arrays: pjrt_array.cc, Line by Line — A jax.Array spread over 512 chips is one Python object, and underneath it is a list of ordinary single-device buffers plus a description of how they tile the whole. This 696-line C++ file is the class that holds those two things together, and the walk reads all of it. The validator that runs before any array exists, which matches buffers to devices strictly by position and compares only the devices this process can address. The six construction paths, one of which quietly skips the validator. The three copy semantics that the header describes as different and that one function implements identically, with the TODO admitting it. Disassembly, which hands back one array per shard without moving a byte. The read back to the host, which refuses anything but a single shard unless the array is replicated. And a copy path where the same invariant is re-derived against a different device list. The signature exhibit runs the validator itself: pick a sharding, perturb one buffer, and watch which of the five checks refuses it and with which string.
• The CPU Thunk Executor: thunk_executor.cc, Line by Line — A GPU gets a driver queue to submit work to. A CPU gets nothing, so the compiled program arrives as a flat list of runtime actions called thunks, and something has to work out what may run at the same time. This 841-line C++ file is that something, read end to end. Each unit declares which buffers it reads and which it writes, and every dependency in the program is derived from those declarations and nothing else, by a conflict test where read-after-read is free. Then a transitive reduction deletes the edges other edges already imply and hands each node a priority equal to how much it unblocks. At run time it is one atomic counter per node and a ready queue: a thunk finishes, decrements its successors, and any counter that hits zero names work that can start now. We read the two edge kinds and why sharing a collective communicator is a weaker constraint than sharing memory; the three thresholds that quietly make a small program single-threaded; the recursive halving that spreads forty ready thunks over four threads without forty queue pushes; and the completion count over sink nodes that makes success and failure end the same way. The signature exhibit is a real ten-thunk graph you step one loop iteration at a time, with all three ready-queue classes.
• The GPU Codegen Path: From a Fused HLO to a Kernel — When the optimizer stops, your model is a few hundred fusions and a handful of library calls, and none of them is machine code yet. This reads the two decisions that finish the job, at a pinned commit of openxla/xla. First the kind: a string an upstream pass stamped on the backend config, or, when there is no stamp, ten questions asked about the fusion roots in a fixed order until one answers yes, with a tenth that always does. Then the class the switch builds from that kind, including the two extra questions the loop arm asks before it settles, one of which turns a copy into a memcpy and skips the kernel entirely. Then the three families behind those classes: six emitters that write a module in an MLIR dialect XLA defined for the purpose and lower it through forty-one passes to LLVM; the Triton path, which builds one module and hands it to another compiler; and the library paths, where cuBLAS arrives as a custom call and cuDNN arrives as a fusion, for reasons that follow from what each library accepts. Then the autotuner: the only pass that compiles and runs real kernels while it is still compiling, the four ways it can get a config, the sentences it uses to refuse, the two-microsecond window that lets a thriftier kernel win over a faster one, the clustering that treats correctness as agreement rather than truth, and the two-tier cache with a key caveat the source states out loud. The signature exhibit runs the real dispatch function over seven candidate fusions and shows which question fires and which class gets built.
• XLA Collectives: One Guest List, Then One Order — A collective is the one instruction in a compiled program that cannot be executed by looking at its own operands, because it needs values from machines that do not share memory with this one. This reads the whole subsystem at a pinned commit of openxla/xla, from both sides of the line. On the compiler side: the list of integers an instruction carries, the two optional fields that decide whether those integers are replica ids, partition ids or positions in a flat enumeration of the device grid, and the switch with four arms that multiplies one written group out against a mesh into the sets of global device numbers the runtime will actually open communicators for. Then the flattened id and the device number, which walk the same grid in opposite directions and disagree on any mesh with more than one of each. On the runtime side: ranks, communicators, cliques and the factory that makes them, all defined without naming a vendor; the future that means launched rather than finished, and means two different things on GPU and CPU; the rendezvous that agrees on membership before any bytes move; and the three separate places where the answer to how do we avoid a deadlock is the same answer. The signature exhibit runs the real group arithmetic and the real verifier checks over a mesh you dial.
• Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot — One keyword argument runs a TPU kernel on a laptop, and almost every confusing Pallas afternoon comes from not knowing what that run checked. This reads the machinery at a pinned commit of jax-ml/jax. There are three interpreters behind one argument, not one: a generic scan over the grid that models the arithmetic and removes memory entirely, a TPU mode that simulates HBM, VMEM, DMAs and semaphores and carries a happens-before race detector, and a real chip, the only rung with a compiler that can refuse you. What each rung catches and what it structurally cannot, drawn as a walk. Which import names which chip, with the memory-space tables side by side and the SMEM collision that means scalar memory on one and shared memory on the other. What interpret=True literally executes: a while loop over the grid, dynamic slices for blocks, outputs and scratch pre-filled with NaN. The padding the interpreter adds before the loop starts, computed live from the real function. Why the grid promises a count and not an order, and the one parameter that randomizes it. Then the case worth committing to: an index map one block past the end, which raises nothing, returns finite numbers, and quietly duplicates half the output. The debugging surface as a picker, one question in and one switch out. And the six-cell notebook shape, including the reset cell that exists because the interpreter keeps its simulated memory after an exception on purpose.
• pallas_call, Line by Line: One Call, Three Calling Conventions — Every Pallas kernel on every backend enters JAX through one 1381-line Python file. This walks it whole, at a pinned commit of jax-ml/jax. The file declares a single primitive, then spends the rest of itself teaching JAX what that primitive means under each transform: an abstract evaluation that decides the output types and refuses a non-manual mesh, a lowering rule that swaps its own body for a scan when you ask for interpret mode, a batching rule that answers vmap by growing the grid rather than looping, and a public function whose real work is turning a pile of keyword arguments into one GridMapping. The signature exhibit runs the real slot arithmetic out of that file: dial the grid, the scalar prefetch operands, the inputs, the outputs and the scratch, and watch the three orderings the primitive, the kernel and the index map each demand.
• The Mosaic TPU Pipeline, Line by Line — A Pallas kernel on a TPU computes from VMEM, a scratchpad of on-chip memory small enough that the array you passed in does not fit. So every block has to be copied in from HBM before the core can touch it, and copied back out afterwards, and if the core waits for either of those the matrix unit sits idle for most of the schedule. One file decides when each copy is issued, when the core is allowed to wait for it, and how many copies of a block VMEM holds at once. This walks that file at a pinned commit of jax-ml/jax: the buffered ref and its four independent cursors, the semaphore per slot, the four boolean predicates that are the entire schedule, the prologue that primes the buffers stage by stage, the loop body whose line ordering is the only reason anything overlaps, and the epilogue that waits for the last send nobody else waited for. The signature exhibit runs the real predicates over a grid you dial and prints the resulting schedule, slot by slot.
• Pallas on a TPU: From a Jaxpr to the Mosaic Dialect — A Pallas kernel is a Python function, but by the time a TPU sees it the Python is long gone: what the backend receives is a jaxpr, and what it emits is an MLIR module in the Mosaic dialect wrapped inside one HLO custom call. This reads that translation at a pinned commit of jax-ml/jax. The type mapping first, where a ref becomes a memref carrying a memory-space attribute and a value becomes a vector, and the four-step chain that turns a BlockSpec with no memory space into the literal string vmem. Then the window parameters that decide whether a block is copied in for you or left in HBM for your own DMA, and the two allocators behind run_scoped, one calling alloca and one calling a semaphore allocator. Then the table: 135 primitives carry a TPU rule at this commit, computed from the real registrations and filterable by core type and by the MLIR dialect each rule builds into. Then the two doors to the MXU, one of them a single tpu.matmul and the other five primitives that name an accumulator and an MXU index. It closes on refusal: the dtypes the lowering will not take, the block shapes it rejects before lowering starts, and the one message that means no rule exists at all.
• BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces — A Pallas kernel never sees your array. It sees one block of it, chosen by a Python function you wrote and a loop the compiler ran, and this page reads the arithmetic that connects the two at a pinned commit of jax-ml/jax. The grid is a loop nest and the kernel is its body, so the kernel runs prod(grid) times. A BlockSpec is a block shape plus an index map, and the index map returns block indices rather than element indices: the start of the slice is the block size times the number your map returned, which is one match statement in core.py and nothing more. From there the page walks the five kinds of block dimension and what each one asks your map to return, the three defaults and why omitting a BlockSpec entirely does not mean the whole array but an unblocked ref in whichever memory the backend chose, the two memory-space vocabularies and where DEFAULT lands on a TPU, and the conversion that turns your closure into a traced jaxpr along with the three rules that conversion enforces. Then the honest parts: a block shape that does not divide the array runs anyway, on a full-sized block, reading padding whose values the documentation tells you to assume are garbage, discarded on output and not on input; the grid runs with the last axis fastest on the interpreter and in an order you may not assume on a chip; and a TPU wants the last two dimensions of your block to be multiples of 8 and 128. The signature exhibit ports the real block-index arithmetic, the ceiling division, the start indices and the padding rule, so you can dial an array shape against a block shape and read the slice the kernel would receive.
• Pallas on a GPU: Two Backends, and What Each One Lets You Write — One pallas_call on an NVIDIA card reaches one of two compilers, and they do not accept the same kernel body. Each backend is, quite literally, a dictionary from JAX primitive to lowering function, and a primitive with no entry raises before a single instruction is emitted. This reads both dictionaries at a pinned commit of jax-ml/jax: 101 primitives on the Triton side, 117 on the Mosaic GPU side, 59 in common, and a matmul that lowers on one and has no rule at all on the other. Then it reads a production kernel library at a pinned commit of google/tokamax to see what a team does with that: a Triton attention kernel in one file, a Mosaic GPU dispatcher that owns no kernel, and two kernel bodies forked by GPU generation because the accumulator moved out of the registers on Blackwell. The signature exhibit checks one softmax body against the real rule tables, line by line.
• Flash Attention on a TPU, Line by Line — The flash attention algorithm fits on a napkin. The Pallas kernel that runs it on a TPU is 1715 lines: four kernel bodies, three separate launches, eleven block sizes, and a four-line diagonal test consulted from ten places. This walks the whole file at a pinned commit of jax-ml/jax.
• Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop — Flash attention treats a causal mask as a value: compute the whole block of logits, then overwrite half of them with negative infinity. Splash attention treats the same mask as a schedule. Before the kernel is ever traced, a Python pass cuts the mask into blocks the size of one grid step, labels each block empty, partial or full, and turns that label grid into three small integer arrays that live in TPU scalar memory. The kernel then reads its own loop structure out of them: an empty block is not computed and its data is not fetched, a full block skips the masking entirely, and for a banded mask the grid itself is rebuilt narrower. The signature exhibit and the dial both run the real classification out of the pinned file, so the block counts on this page are the counts the compiler would produce.
• The tokamax Op and Its Autotuner, Line by Line — A fast kernel is never one kernel. It is a family of implementations and a family of tilings, and something has to choose between them on every call. Tokamax, the fused-kernel library Google builds on Pallas, answers that in one 701-line Python file. This walks it whole at a pinned commit: the frozen Op dataclass and the two hooks a backend overrides, the twelve lines that turn a call into a dispatch key made only of argument shapes and dtypes, the five-step ladder that decides whether you get an explicit config, a cached measurement, an unmeasured guess or an error, the measurement loop that fills the cache in, and the payload the file writes into the compiled program so a captured production trace can be tuned on a machine that never runs your model. The signature exhibit is that real key, opened against the library’s own shipped cache: pick a ragged-dot workload and a chip, and read every cached implementation with its winning tiling and its measured median.
• Attention for sm90, Line by Line — A production flash attention kernel for the H100, read whole at a pinned commit of google/tokamax. The algorithm is the one everybody knows; what a real kernel adds on top of it is six hundred lines of scheduling. One thread block launches three warpgroups and each branches immediately into a different role: two do the arithmetic on 232 registers a thread, and the third drops to 40 and does nothing but issue transfers, running a stage ahead of the maths. Automatic synchronisation is switched off at launch, so every handshake in the file is one of sixteen numbered hardware barriers placed by hand, and running out of them is a real failure mode. The page walks the constants, the shared-memory plan, both configuration heuristics, both warpgroup roles and the launch, then steps through one query tile: which slot each warpgroup is holding, how far ahead the loader has run, and where the causal mask starts biting. A second exhibit runs the file own shared-memory estimate so you can watch a tiling stop fitting in 227 KiB.
• 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.
• Running a Tunix Recipe: From a YAML File to an RLCluster — There is no tunix command. You run a module, hand it a YAML file as the first positional argument, and a few hundred lines of config code turn that file into a tokenizer, a dataset, a device mesh per role, three or four models, a rollout config, a cluster and a learner, in that order, before a single token is generated. This reads the whole launch surface at a pinned commit of google/tunix: the four sources config is merged from and the exact rule that makes setting one key by both an environment variable and the command line an error; the two base YAML files and the eleven keys that exist in only one of them; the five sections where a partial override silently deletes every sibling key you did not restate; the unknown-key check that catches a typo at the top level and cannot see one a single level down. A recipe turns out not to be a YAML file at all but a Python module exposing create_dataset, and a reward function is any module-level function in a file you name. The signature exhibit walks one real launch command from the file to the objects, with the arithmetic the code actually does at each step.
• rl_cluster.py, Line by Line — Reinforcement learning on a language model means running several copies of that model at once, each doing a different job: one samples completions, one scores them, one is a frozen snapshot you measure drift against, and one is taking the gradient steps. This file is the object that owns all five roles and decides where each of them physically lives. The answer is one dictionary from role to mesh, compared for object identity in exactly two places, and everything else follows: whether the sampler is a second copy of the policy or a second name for the same buffer, whether a weight sync moves bytes across the interconnect or nothing at all, and whether a LoRA reference model costs any memory. We read all 1,252 lines of the construction and step path at a pinned commit, then dial the three settings that decide placement and watch the file’s own branches fire.
• The RL Learner Loop, Line by Line — One file in google/tunix, 834 lines, read whole. It never touches a model and never computes a gradient. What it does is arithmetic on batch sizes and a handoff between two threads: a producer that pulls prompts, generates completions and computes advantages, and a trainer that pulls finished examples off a queue and hands each one to the optimizer. The page derives the six batch sizes from the config, walks the producer and the consumer region by region, and ends on a computed step timeline where you set the sizes and watch rollouts and gradient steps either overlap or take turns, depending on one line that compares two meshes.
• The GRPO Learner, Line by Line — Six hundred lines of google/tunix that add group-relative policy optimization to a training loop they do not contain, and that compute neither the advantage nor the loss the file is named for. Both arrive as function pointers, fetched from a process-global table by a string that lives on the config, filled by an import the file never mentions again. The page walks the whole file at a pinned commit: the config whose docstring documents a field it does not declare, the constructor that hands the actor trainer its loss, the one method that turns prompts into a differentiable record, the seventy lines that measure how far the rollout sampler has drifted from the trainer, and the four abstract slots that are the entire delta over the base learner. Then a group of four completions you score yourself, with the real estimator and the real clipped loss ported line for line, and the tied group that costs a full step and teaches nothing.
• 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.
• The Tunix Sampler, Line by Line: Generation as Two Compiled Functions — Every RL post-training step needs completions, and in Tunix the in-tree sampler is what produces them. It is one file, and its whole design falls out of one constraint: JAX wants fixed shapes. So the prompt gets left-padded to a power of two, the output lives in a buffer sized before the first token is generated, the KV cache is allocated once and donated to the compiler, and the decode loop is a lax.while_loop that runs until the slowest sequence in the batch stops. We read tunix/generate/sampler.py top to bottom at one pinned commit, then compute the decode timeline and the cache filling for a batch you dial yourself.
• Rollout Backends and Weight Sync: Getting New Weights Into a Sampler — An RL loop trains one copy of the policy and samples from another, so every batch ends with a call that moves the updated weights across a device mesh, renames them for whichever inference engine is on the far side, and quietly does nothing at all when the two already agree. We read Google's tunix rollout backends and its reshard function at one pinned commit, then compute the reshard plan for a real Llama-3.1-8B parameter tree: which arrays cross, how each is cut on the other side, and what a batch boundary actually costs.
• The Trajectory Collect Engine, Line by Line — A normal RL rollout is one prompt in, one completion out. An agentic rollout is a conversation: the model writes a tool call, an environment runs it, the result comes back as a message, and the model writes again, for however many turns the task takes. This walk reads the 763-line file in Google’s tunix that owns exactly one of those episodes: how it counts steps, where the only deadline in the file actually sits, which tokens end up carrying gradient, and how dozens of these run at once and stream their finished trajectories to a learner that never stops training.
• The Tunix PEFT Trainer, Line by Line — A LoRA fine-tune freezes the base model by leaving it out. This walk reads tunix/sft/peft_trainer.py at one pinned commit: the type tag that decides which parameters an optimizer is even built over, the two update paths inside one train step, the gradient accumulator that carries a denominator, and the checkpoint that writes only the adapter.
• 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.