The library

Every visual explainer on Rudrite Explainers, filterable by kind and field.

• The System Design Cheat Sheet — A hundred design prompts look like a hundred things to memorize — they are one method re-dialed. The meta layer above the worked examples: the eight-move arc of the round (scope, size, break the naive version, deliberate, go deep in dependency order, sweep the failures, wrap); the numbers you carry in (the Jeff Dean latency ladder on a live log scale, per-box capacity rules of thumb, the four QPS/storage/bandwidth formulas); the building-block menu with the failure an interviewer probes the moment you place each one; the six decision tables (SQL vs NoSQL, strong vs eventual, cache-aside vs write-through, push vs pull, sync vs async, precompute vs on-demand), verdict-first; and the six cross-cutting invariants a senior never skips (idempotency, backpressure, hot keys, fan-out, honest exactly-once, checkpointed resume). The signature exhibit is a prompt-decoder: pick a classic prompt and watch which blocks it pulls and which invariant dominates.
• The Audit-Proof Ledger — An auditor asks three questions a mutable database cannot answer: what was the balance last March, is this history real, and did your logic still hold after last week’s deploy? The answer is to stop storing the balance and start storing the events — every validated fact, in order, forever. Built from zero: commands versus events (intent versus recorded fact), the deterministic state machine whose replay is byte-for-byte reproducible, why the same events always yield the same balances, CQRS read models rebuilt per consumer, and the systems engineering that makes it fast (append-only files, memory-mapping, snapshots to bound replay) and reliable (a Raft-replicated event log where the leader appends and followers replay). The signature exhibit replays the same twelve-event log twice for identical balances, then corrupts one historical event and watches the divergence surface at exactly that event — computed live from a real reducer.
• Design a Trading Dashboard — A million people watching prices that move millions of times a second, and a browser that can render maybe ten of them. Built from zero: the read-side scope (the display, not the matching engine), a commit-first envelope where naively forwarding every tick is ≈240 GB/s and conflation cuts it to ≈4.8, market-data fan-in and normalization across venues, then the signature idea — conflation, keeping one latest slot per symbol and emitting ~4 times a second (exactly what Interactive Brokers ships), pub/sub fan-out over a WebSocket gateway fleet, snapshot-then-delta on subscribe with sequence-gap resync, per-client backpressure that conflates again at the socket, tiered update rates that follow attention, OHLC candle roll-ups computed on the full tick stream, portfolio P&L as the live price joined onto positions, and the two-lane guarantee split — lossy display, exact books — that runs through the whole design.
• Design a Chat System — Two people, a message, and the expectation that it arrives instantly, in order, and never vanishes — even when the recipient is in a tunnel. Built from one polling request up to a stateful WebSocket fleet over a durable inbox: the commit-first envelope that sizes everything by open connections (not requests), the honest transport fork (short poll vs long poll vs WebSocket), a session registry that finds a user across the fleet, store-then-notify with per-recipient inbox queues, Snowflake IDs because created_at collides, a bucketed key-value store for history because reads are long-tailed, the signature offline-delivery exhibit where a per-device cursor replays exactly the gap, presence via heartbeat + flap-suppression + pub/sub fanout, and a failure sweep whose surprising box is a chat server dying with all its connections.
• Design a Video Platform — Someone uploads a 4K clip; seconds later millions can stream it on a train, a laptop, a TV — and the bill to deliver it dwarfs the bill to store it. Built from zero: a commit-first envelope on the CDN egress bill (5M users × 5 plays × 300 MB is 7.5 PB a day, and the number should scare you), why one server melts, the two-lane upload with pre-signed URLs writing straight to storage and GOP-aligned resumable segments, the heart — transcoding as a DAG of tasks where every segment × resolution encodes in parallel and a dead worker costs one segment not one video, adaptive bitrate over a per-title ladder with the player choosing rung by rung, CDN economics and the long tail (cache the hot set, serve the cold tail from origin, lazily transcode what nobody watches) — then the failure sweep whose sharpest box is a poison video that turns your own idempotent retry into a fleet-killer.
• Design a Payment System — A million transactions a day is only about ten a second — so the hard part was never throughput, it was never losing or duplicating a single cent. Built from zero: the hosted-payment-page flow that keeps card numbers off your servers entirely (nonce → token → charge → webhook, with the PCI reality stated plainly), idempotency keys that survive both a double-click and a payment-processor retry, the double-entry ledger where every movement is a debit and a credit that sum to zero, why money is integers and never floats, a retry queue with a dead-letter tail for the charges that fail, and reconciliation as the last line of defense — the nightly diff against the bank’s settlement file and the three classes of mismatch it turns up. The signature exhibit sends one payment across the whole system with its idempotency key visible, replays it to watch every stage no-op, then breaks the ledger write and lets reconciliation catch the gap the next morning.
• Design an Email Service — A billion users, forty messages a day each, kept forever — email is a storage problem wearing a messaging problem’s clothes, and almost none of it is under your control once a message leaves the building. Built from zero: a commit-first envelope that lands on the yearly petabytes, the send and receive pipelines threaded through queues so a slow recipient never blocks anyone, the mailbox data model at the heart — partition by user, folders, time-ordered message IDs, and the honest read/unread denormalization a wide-column store forces on you (two tables, move the row) — the search index as its own write-heavy LSM, and deliverability as its own discipline: IP reputation, dedicated-IP warm-up, SPF/DKIM/DMARC in a breath each, and the bounce/complaint feedback loops — then the failure sweep whose sharpest box is your own IP landing on a blocklist.
• Design a Notification System — One event — “your driver is arriving” — has to reach a phone, and you own almost none of the road: the last mile belongs to Apple, Google, Twilio, and an email provider, and none of them promises the message arrives. Built from zero: a commit-first envelope on 10M push + 1M SMS + 5M email a day, why one inline sender collapses, the fan-out core where a single event lands in one queue per channel so a Twilio outage never blocks a push, the device-token table and the templates / user-settings / rate-cap layers that keep you from spamming a real person, the reliability spine (persist-then-queue, retry with backoff, and the honest at-least-once truth that dedupe reduces but cannot erase), the open/click tracking loop — then the failure sweep whose quietest box is a compliance bug that ships a notification to someone who opted out.
• Build a Report and Fan It Out to Millions — One prompt hiding two systems: a periodic batch pipeline that computes 20 million personalized reports, joined to a delivery ramp that mails every one inside a 3-hour window over a provider you don’t own. Built from zero: the scope split (compute vs deliver), a commit-first envelope where a ~2,000/s send ceiling barely clears the deadline, the naive cron loop and why it dies, the compute DAG with a pinned snapshot and per-task retries, the central precompute-vs-render-on-open deliberation tuned to the open rate, rendering to immutable artifacts behind short-lived signed links, then the delivery half — batching, a token-bucket rate shaper clamped to the provider ceiling, idempotency keys so nobody gets the report twice, suppression before the send, a checkpointed cursor so a crash resumes instead of restarting, and a live completion burn-down that proves all N went out by when. Distinct from a notification system (linked, not duplicated): this is the report COMPUTE joined to the delivery ramp.
• Design a Proximity Service — Open a maps app, tap "restaurants near me," and a few of the 200 million businesses on Earth come back in under a second. This is the retrieval shape under Yelp, "find nearby drivers," and every store locator — built from zero: why a plain index on latitude and longitude quietly falls apart, geohash derived by hand (recursive halving, prefix as zoom, base-32) with the two boundary traps that make a naive query silently miss results and the eight-neighbor fix that catches them, an honest geohash-vs-quadtree-vs-S2 deliberation, the compound-row data model, the read path and why the user’s own coordinates are a terrible cache key — then "when the points move": nearby-friends over a pub/sub channel per cell and a WebSocket fleet, honest about the write volume it costs.
• Top-K & Heavy Hitters — Who are the top ten right now? Easy at a thousand rows, brutal at a billion. Two questions wear one name: a live leaderboard of 25M players, and the heavy hitters of a firehose you can never store. Built from zero — why SQL ORDER BY dies, how a Redis sorted set answers rank in O(log n) with a skip list, how the board shards and why hashing breaks the rank query, then the streaming half: a count-min sketch that counts a billion keys in a kilobyte (and sometimes lies), a min-heap of the top ten, and lambda reconciliation when the number has to be billing-exact.
• 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.
• Design a Transaction Log on Object Storage — A cloud bucket can swap one file atomically and never many at once — so how does a pile of Parquet files on S3 become an ACID table? The Delta-shaped question, built up from first principles: the naive folder and the three ways it dies, the write-ahead log as the source of truth, the put-if-absent commit (and why S3 needs a tiny coordinator), optimistic concurrency with the conflict rules raced live, checkpoints against a computed cold-read cost, compaction, time travel, and the log doubling as a message queue — then an honest Delta vs Iceberg vs Hudi.
• Design a Matching Engine — A stock exchange makes four promises at once — fairness, determinism, microsecond latency, and never losing an order — and one data structure keeps all four. Built from zero: what an exchange must guarantee, the order book as sorted price levels each holding a doubly-linked FIFO plus an id-map (place, match, and cancel all O(1), and the singly-linked trap that ruins it), a live book where you fire a market buy and watch it sweep the ask side with honest fills, the sequencer that stamps every order so the same input always replays the same fills, the single-box mmap event bus that beats the network by three orders of magnitude, hot-warm failover, and the p99 tail that a garbage-collection pause quietly ruins.
• Design a Sandboxed Code-Execution Service — Run code an LLM wrote seconds ago, for thousands of mutually suspicious tenants, without ever letting one read another’s files or take the host down — the isolation spectrum from a shared process to a hardware-walled microVM, warm pools and snapshot/restore, cgroup limits, and network policy, priced in real milliseconds and megabytes from the Firecracker and gVisor docs.
• 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.
• How Google Spanner Works — A database that spans the planet and still behaves as if every transaction ran one at a time. The trick is TrueTime — a clock that answers with a range instead of an instant, so the database can wait out its own uncertainty. We walk how the data is split and copied, the waiting rule that keeps ordering honest, and when you do not need any 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.
• A Write-Ahead Log You Can Implement in an Hour — The low-level-design round asks you to make a key-value store survive a crash — and the whole answer is one small file. Append a length + CRC32 framed record before you touch memory, fsync on the durability dial you choose, and on restart replay the log and drop the torn tail. We author the ~120-line Python module, crash it in a child process, and walk it top to bottom.
• The GIL, Honestly — One mutex around the interpreter lets only one thread run Python bytecode at a time — and yet x += 1 can still lose an update. Watch the race survive the GIL op by op, then pick threading, multiprocessing, or asyncio from a computed wall-clock, and see what 3.13 free-threading really changes.
• Design a Stream Processor with Exactly-Once — How a system that counts events as they stream past can crash, recover, and never count anything twice: periodic snapshots of its memory, an honest way to handle events that arrive late, and an output stage that only commits when everything agrees — drawn, computed, and animated.
• Design a Distributed Cache — Start with one cache box in front of the database; end with a fleet that survives a celebrity post pulling 300k requests a second at a single key. How keys spread across machines, what gets evicted when memory fills, and the defences for when one key goes nuclear — drawn and computed.
• Back-of-the-Envelope: the Numbers That Design Systems — Before you draw a box, you do the arithmetic. Powers of two, the latency ladder, and a live calculator that turns daily users and payload size into QPS, storage, bandwidth, and cache — every formula shown.
• Threads, Locks & the Anatomy of a Race — Share memory between two threads and the number of ways their steps can interleave explodes — most correct, a few catastrophic. Walk the canonical lost update, scrub an adversarial schedule op by op, then hold the invariant with a mutex, condition variables done right, and a deadlock you can watch form.
• Two Writers, One Row — The last hotel room, two customers, both screens say “1 available,” both click book. Walk the interleaving that double-books, then hold the count with three defenses — a pessimistic lock, an optimistic version check, a database constraint — and when the row spans two services, a saga.
• Consistency, Quorums & CAP — Replication forces copies to disagree. The consistency ladder, the N/W/R quorum rule (W+R>N), and CAP stated precisely — the real choice is C-vs-A only during a partition — turned live on a 5-replica cluster.
• How to Design a System in 60 Minutes — A system design interview is not a memory test — it is four phases on a clock. Restate the problem, commit to one design, go deep where it matters, then break it on purpose. The framework, drawn as an interactive script you scrub minute by minute.
• Preparing for the Agent-Infrastructure Interview — A design invite lands with a rubric axis you have not rehearsed for — "Automation Depth & AI Leverage" — sitting right next to the familiar platform-design one. What that new line actually grades, the four question territories it draws from, and how to walk in ready on both.
• Rate Limiting: Four Algorithms, Honestly Compared — Token bucket, leaky bucket, fixed window, sliding window — one bursty arrival pattern replayed through all four at once, every verdict computed from the real algorithm. Where each one lets a burst through, and what it costs to remember the past.
• Idempotency & the Exactly-Once Illusion — Over a lossy wire, “did it happen?” is unanswerable — so retries are the default reality. Exactly-once delivery is impossible; exactly-once effects are not. The property that buys them back is idempotency, drawn and computed.
• Design Ad-Click Aggregation (Lambda vs Kappa) — Count a billion ad clicks a day two ways at once — a dashboard fresh within seconds, and a total accurate enough to bill against. How duplicate clicks are weeded out, how late-arriving clicks are handled honestly, and why one replayable pipeline can do the work of two.
• Design a DAG Job Scheduler — A scheduler that keeps its plans in memory is a demo: the moment it crashes, jobs are lost — and running two copies for safety makes jobs run twice. Build up to schedulers that share one database safely, so no job is ever lost or double-run — drawn and animated.
• Unique IDs at Scale — A single auto-increment counter is trivial — until you shard and it becomes the hardest thing in the system. Watch the fix: pack a clock, a machine id, and a per-millisecond counter into 64 bits, with a scrubbable Snowflake bit-strip and clock-skew you can break — then watch the same ID become a base-62 short URL.
• From One Server to Millions of Users — The scaling journey as a playable timeline. Start with everything on one box, then let each real bottleneck force the next move — split the database, load-balance a stateless web tier, add read replicas, cache, autoscale across zones, shard, and decouple write bursts with queues.
• Consistent Hashing — Add a machine to a sharded system and plain hash-mod-N reshuffles almost everything. Consistent hashing puts the nodes on the same ring as the keys — walk clockwise to the owner — so a membership change moves only one arc, and virtual nodes even out the load. Drawn, computed, and animated.
• The Shuffle: How a Cluster Moves a Join — A join is trivial once the matching rows share a machine. The shuffle is the all-to-all network move that gets them there — repartition every row by a hash of its key, and everything hard follows: why it dominates query time, why one hot key strands a straggler while the fleet idles, and how adaptive splits, salting, and broadcast joins break it.
• Design a Multi-Tenant Query Engine on Object Storage — Thousands of customers, ten thousand queries at once, sub-second dashboards next to multi-hour batch jobs — all on discounted machines the cloud can take back at any moment. Split the coordinating brain from the number-crunching muscle, schedule fairly, and make a query survive its own workers being yanked mid-flight.
• Design a Durable Key-Value Store — Start with a lookup table in one server's memory; end with a storage engine that survives crashes and power loss. One idea — a log file written before anything else moves — is the whole story, drawn and computed.
• Design a Metrics & Monitoring System — Why billions of measurements a day outgrow an ordinary database and need a purpose-built one — the compression tricks that make the data fit, the label explosion that can sink the whole system, and an alerting pipeline that watches itself.
• Design S3-Like Object Storage — From one server writing files on disk to a store holding a million terabytes that survives an entire data centre dying — the bytes cut into recoverable pieces, spread across buildings, and repaired faster than they fail.
• Design Search Autocomplete — The suggestions that drop down as you type — built from zero. A trie (a tree keyed by letters) with the top few completions pre-computed and cached at every node, so a keystroke is answered by reading a list, not by searching; why a live database query per keystroke dies at ~20× the search rate; the offline pipeline that turns query logs into a fresh trie snapshot; the sub-100 ms serving path; sharding by prefix and the fix for the uneven letters; and the failure sweep whose quietest box is a trending query that stays invisible for a week.
• Design a Web Crawler — The bulk downloader behind a search engine — built from zero. Why naive breadth-first crawling is rude, and how the URL frontier makes it polite: front queues that prioritise (by PageRank, traffic, freshness) feeding back queues that enforce one host per queue, one worker, a delay between fetches. Plus robots.txt and its cache, the two dedup structures (a bloom filter of seen URLs, a hash of seen content), DNS caching and the blocking-resolver trap, spider-trap defenses, and a freshness strategy — with a live frontier you can flood from one host to watch politeness throttle it while the others proceed.
• Design a Distributed Message Queue — How services hand work to each other through a queue that never loses a message. Start with one sender and one receiver; grow it until it survives crashes and slow consumers — drawn and animated.
• Design a News Feed — Open the app and see a fresh, ranked stream of posts from everyone you follow. The whole design is one fork — do you push each post into every follower’s feed as it’s written, pull them all together at read time, or split the difference? Built from zero: a commit-first envelope on 10M users, the push/pull/hybrid deliberation with cost meters computed live, the celebrity hot-key that melts the write path, the feed cache kept as bare <post-id, user-id> lists and why not full posts, the hydration step, an honest paragraph on the ranking hand-off, and the failure sweep.
• Design a File-Sync Service — Drop a file in a folder and it appears, seconds later, on every device you own. Build the Dropbox shape from zero: cut files into content-hashed 4MB blocks so editing one block ships one block, guard a tiny metadata index with strong consistency while the bulky bytes stay loose, nudge devices with a held long-poll connection instead of a poll storm, and reconcile offline edits into first-writer-wins conflicts that never silently merge — with the delta-sync bandwidth save computed live.
• Design Google Maps — Two systems wearing the same skin: a pyramid of pre-drawn tiles you pan and zoom, and a road graph cut into tiles you route across. From the naive one-server cut to a static tile pyramid on a CDN, hierarchical routing tiles searched with a live A*, ~1M-QPS location ingestion, and live-traffic ETAs — 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.
• CPython’s queue.py, Line by Line — The real bounded blocking queue from the standard library you already import. One mutex, three Conditions, and the not-empty / not-full invariants — walk the whole file top to bottom and watch the monitor block, wake, and drain.
• How Redux createStore Works, Line by Line — Redux’s entire state engine is one function that returns a closure. Walk the real createStore.ts source top to bottom — dispatch, subscribe, the isDispatching guard, and the INIT that boots your state tree.
• functools.lru_cache, Line by Line — The canonical interview prompt — build a thread-safe LRU cache with O(1) get and put — answered in production Python. Walk the real functools.py: the flattened key, the circular doubly linked list, the pointer surgery that moves a hit to the front, the reuse-the-oldest-node eviction trick, and the one RLock that keeps it all consistent.
• Python’s Dict, Under the Hood — A namespace, a set of keyword arguments, and an object’s attributes are all the same thing: a dict. Walk the CPython hash table that powers them — open addressing, the perturb probe sequence, the compact two-array layout that made insertion order a language guarantee, resizing, and SipHash — with the probe walked live over a real hash.
• Python Gotchas That Fail Interviews — Eight snippets that look obvious and aren’t — mutable default arguments, late-binding closures, is vs ==, chained comparisons, mutation during iteration, shallow copies, shared class attributes, and a finally that eats a return. Predict each output, then meet the one evaluation-model rule underneath. Every result run in real CPython, every claim traced to the docs.
• Iterators, Generators & the Dunder Protocols — Python has no interfaces — a syntax works on your object when the object supplies the right dunder. Desugar a for-loop into iter/next/StopIteration, watch a generator freeze mid-frame on yield, and tour the protocol family behind in, len, bool, with, and ==. Computed and animated.
• Errors, Edge Cases & Testing in the Room — Juniors code the happy path; seniors enumerate what breaks first. The edge-case checklist as a habit, raise-with-context and exception chaining, assert vs exceptions and the -O caveat, a five-minute inline self-test, and property thinking without a framework — worked on one binary search, every snippet run in real Python.
• Designing a Clean Python API Under Pressure — The low-level-design round grades your interface, not your diagram. Build one small key-value store live — but design only: derive the API from the call site first, state and enforce the invariants, shape the exception surface, add dunders where they earn their keep, and let the docstring be the spec. Every snippet runs in real CPython.
• Idiomatic Python — The same program, written twice — once with a C accent, once the way Python wants. Comprehensions, unpacking, the right dict idiom, dataclasses, context managers, EAFP, f-strings, pathlib, and truthiness — each as a before/after pair, with the honest case for when NOT to reach for the idiom. Every snippet run in real CPython.
• Writing Python Interviewers Trust — One real interview problem — parse a log file, report the top-N error codes per service — solved twice and walked side by side: the rushed 26-line function (single-letter names, a bare except, prints, mutates its input) versus the trustworthy version (small functions, guard clauses, honest exceptions, a pure core and a thin I/O edge). The refactor as six trust signals you can name and practise.
• 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.
• Design a Webhook & Trigger Delivery Platform — Thousands of customer servers, each unpredictable — slow, offline, flaky at 3am — and one bad one must never touch the rest. We compute the fan-out arithmetic (events → deliveries, retry amplification when 5% of endpoints are down, storage for a 3-day retry horizon), watch a naive POST-in-the-request-loop stall on a single slow endpoint, then build the real design one mechanism at a time: a durable queue off the request path, HMAC signature verification in both directions, per-endpoint token buckets, a jittered retry ladder to a dead-letter queue, honest at-least-once semantics, and ordering that costs only the customer who asked for it — closing on Stripe, GitHub, and Svix's own published numbers.
• The OOD Interview, Decoded — The object-design round looks like it wants pattern trivia. It is really grading one thing: can you turn four fuzzy sentences into a handful of objects with clean seams? We make the method legible — the nouns/verbs pass that turns a prompt into candidate classes and methods, composition over inheritance stated honestly, the invariants you say out loud, the four patterns that actually show up in the room (factory, strategy, observer, state) each with a five-line code moment, SOLID trimmed to the letters that earn their keep, and the secret weapon: most OOD classics are state machines wearing a costume. Click the nouns in a live requirements pass, level a design against the classic failures, and leave with the 45-minute LLD script.
• Three OOD Classics, Worked — The method from the OOD interview page, run on the three prompts that show up most. The parking lot end to end — requirements, the nouns pass, the class model, and the three decisions that actually get graded: spot assignment as a strategy, the ticket as the identity object, fee calculation as a swappable strategy, with the real compact code. Then the reveal that collapses half the catalog: the vending machine and the elevator are the same problem — state machines with states, events, guards, and actions. Drive one live: insert coins, select, dispense; request floors up and down; watch the legal transitions fire and the illegal ones bounce with the guard named. Finish with the missing-guard bugs interviewers probe, and the move that spots a state machine in any prompt.
• The Behavioral Interview, Decoded — The behavioral round feels like a chat and grades like an exam. Two questions run underneath every story you tell — do you fit this team and company, and at what level do you operate — and neither is ever asked out loud. We make both legible: the fit axes a company reveals in its own words, and the four dimensions (scope, contribution, impact, difficulty) that quietly set your level. Dial one true story across five levels and watch the meters move; level a story yourself before the reveal; then sweep the five ways a story gets mis-leveled and sinks.
• Stories That Carry Signal — An interviewer keeps notes, not a transcript — so the real test of every sentence is whether they would write it down. We build the story that survives that notebook: the headline delivered first (problem plus your key action plus why it mattered, in one breath), a core of two or three moments each moving context to thought to action to a result with numbers, and the prepared depth you stock but never volunteer. An honest read on where STAR helps and where it buries the behavior, the red flags that sink an answer — then the noise-pruner: toggle a realistic answer sentence by sentence and watch the interviewer’s notebook fill with only what carries signal.
• The Story Bank — You don’t memorize forty behavioral answers — you build six to eight true stories that between them cover every question a panel can throw, and rehearse until the right one comes to hand on reflex. A drill surface: the nine competency families as attempt-before-peek cards (how each is asked, what the interviewer greps for, the pitfalls, a fill-in scaffold), the six direct questions with different rules, a commit-first drill on the first sentence out of your mouth, and the signature — a live coverage matrix where you map your stories against the nine families and watch it flag the gaps and the one story doing too much.
• Buy the Cheapest Book — It reads like a database query — sort by price, show the top row — but the two words the query hides are the whole design. “Cheapest” is never the sticker: it is a freshly-computed landed cost (item + shipping + tax, in stock, in the right condition) over prices you don’t control. And “buy” is a promise the price has to survive one click later. Built from zero: the scope dialogue that pulls both words apart, the commit-first envelope whose scary number is 11.6k external fetches a second to keep a billion offers a day fresh (which is why you don’t), the naive live-fan-out that waits on the slowest of 33 sellers and its break, the two ingest lanes (marketplace writes you own vs external feeds you crawl), tiered freshness and a crawl budget, the signature landed-cost ranking computed live where the sticker-cheapest almost never wins, a cached aggregate with a staleness clock, search over the catalog, and the load-bearing correctness beat — checkout-time price revalidation with a tolerance policy, reservation against the last-copy race, and an idempotent order that never double-charges. Grounded in Amazon’s Featured Offer and Google Merchant Center docs.
• 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.