JAX is NumPy plus grad, jit and vmap: automatic differentiation, XLA compilation and automatic batching, built as transformations of pure functions. There are no classes to mutate, no global seed to set, no .backward() to call. The training loop becomes a function from params to params — and the same code scales from this laptop to 2,048 TPUs.
jnp.dot, jnp.zeros, broadcasting and slicing all read like NumPy — but every array is traceable. grad turns a function into its derivative, jit turns it into compiled XLA code, vmap turns a one-example function into a batch function. The original function never changes.
f(params, x) → grad · jit · vmapNo nn.Module, no .grad attribute, no optimizer.step(). Parameters are an immutable pytree you pass in; the step returns a new one. a[0] = 5 raises a TypeError; a = a.at[0].set(5) returns a new array. Optax transforms the gradients and jax.tree.map applies the update — state is just data.
{w, b} in → {w′, b′} outEvery random draw takes an explicit PRNG key — PRNGKey(42) is literally the pair [0, 42]. key1, key2 = random.split(key) gives two children; reuse one key and the second “random” draw replays the first exactly. Split per consumer and per step, and results are reproducible across devices and compilations.
split(key) → k1 ≠ k2 · never draw twice from one keyJAX is a compiler for pure functions: hand it a function of explicit arguments and it hands the same function back differentiated (grad), batched (vmap) or compiled (jit) — the state lives in the arguments, the randomness lives in keys, and nothing is ever mutated in place.
By the end you will be able to write a pure-function MLP and train it with value_and_grad and Optax; explain what the first jit call actually does and why a changed batch shape costs a recompile; lift a single-example function to a batch with vmap and read in_axes; see why per-example gradients fall out for free; split PRNG keys correctly and explain the reuse bug; and translate any line of a PyTorch training loop into its JAX equivalent.
You already know the PyTorch rhythm: define an nn.Module, call .backward(), step the optimizer. It works — millions of people use it. But that style has a constraint baked into its DNA: every operation is dispatched eagerly, one at a time, from Python. At a few million parameters nobody notices. At hundreds of billions across thousands of chips, the overhead is the whole problem.
What “eager” costs. In eager PyTorch, every tensor + tensor is a separate kernel launch, and every training step re-interprets the same Python code. The math is fast; the dispatching is not. Here is the arithmetic, kept deliberately rough so you can check it: a large training step is on the order of 10⁹ elementary operations. If each paid even 1 µs of Python-and-dispatch overhead — optimistic for a Python loop — that is 10³ seconds of overhead per step before the accelerator does any math. Fusing the same computation into a few hundred XLA kernels makes the overhead proportional to the number of kernels instead of the number of operations. (This is an order-of-magnitude illustration, not a benchmark: real frameworks overlap dispatch with compute, and launch costs vary by operation. The direction of the effect is exactly this.)
This is why the largest training runs on Earth stopped using eager loops. Google DeepMind trains Gemini on JAX. Anthropic trained Claude on JAX. These are among the biggest neural network training runs ever attempted, and they chose a framework whose defining move is to treat the training loop as a compilable program rather than a sequence of Python calls.
JAX is NumPy with three superpowers. The surface is familiar — jnp.array, jnp.dot, broadcasting, slicing, the same function names and the same semantics. Underneath, every array is traceable, and three transformations become available to any pure function you write:
The transformation names are the verbs you will meet in this lesson. grad differentiates a function. jit compiles it, once, to machine code. vmap lifts it from one example to a whole batch. They compose, in any order, because they all operate on the same thing: a pure function.
The board below is the fastest way to see what “pure function” buys you. Each row is one concept of a training loop written in both dialects — the PyTorch version mutates state inside objects; the JAX version takes state as an argument and returns the new value.
Every row is the same computation written the way each framework wants it. Read the pairs top to bottom to learn the translation — model, parameters, backward, update, batch, devices, randomness, loop.
The whole update is one composed, compiled function from state to state. Notice what disappeared: zero_grad, backward, step.
The source’s honest summary sits underneath this board: use PyTorch unless you have a specific reason — TPU access, per-example gradients, training across hundreds of chips, or a job at the labs that build JAX.
PyTorch stores state inside objects — self.linear holds the weights, optimizer.step() mutates them. JAX refuses both the object and the mutation. That refusal is not taste; it is the price of admission to grad, jit and vmap.
The source puts the two styles side by side. PyTorch hides state:
JAX passes it: the whole model is a function whose weights are an argument. The parameters go in; the prediction comes out. Nothing is stored, nothing is mutated, and the function is testable, composable and compilable precisely because it has no hidden inputs.
Immutability: the concrete difference. NumPy lets you scribble over an array in place. JAX arrays are immutable — this is the first line every PyTorch user gets wrong:
The .at[...] API is a small functional language of its own: .set, .add, .mul, .get, and friends. One subtlety catches out everyone: the update only exists if you keep it. Writing a.at[0].set(5.0) as a statement throws the new array away; the idiom is a = a.at[0].set(5.0). And no, this is not a copying tax: the compiler sees the whole computation and can update a buffer in place whenever it proves nothing else reads it. The restriction is on the language you write, not on the memory the compiler uses.
Pytrees: the universal container. JAX calls nested combinations of lists, tuples, dicts and arrays pytrees. Model parameters are just a pytree — and every transformation knows how to walk one. That is how a whole optimizer update is a single function call over both trees:
Left button: NumPy muscle memory — the exact line every PyTorch user types first. Right buttons: the functional form. Watch the “previous array” line to see what survives each call.
JAX arrays behave like values, not like mutable buffers. The .at[...] indexing API returns a new array; the original is untouched, so the only way to keep the change is to assign it — a = a.at[0].set(5.0).
The error text is paraphrased from JAX’s real message; the behavior — immutable arrays, functional .at updates — is exact. Under the hood XLA can still update a buffer in place when it proves nothing else reads it, so this is a restriction on the language you write, not on the memory the compiler uses.
Try a[0] = 5.0 after an .at update: it fails the same way every time. Immutability is not a mode you leave — it is the data model, and it is what makes the three verbs composable.
A two-leaf tree is enough to check the mechanics. Take params = {'w': 2.0, 'b': 1.0} and grads = {'w': 0.5, 'b': 0.25}, with lr = 0.1. jax.tree.map(f, params, grads) applies f to each matching leaf pair:
That is the entire optimizer abstraction in one line. Adam, SGD, clipping and schedules are each just a different function passed to the same tree walk — which is exactly what Optax packages.
After b = a.at[0].set(5.0), which statement is true?
.backward().PyTorch attaches gradients to tensors — every weight carries a .grad slot. JAX attaches them to functions: jax.grad(f) is a function that evaluates f′. The gradient becomes a value you can call, compose, compile, and batch like any other.
The whole API is one line:
Read the numbers out loud, because this is the source’s own demo. At x = 2, the function x³ has value 2³ = 8. Its slope there is 3·2² = 12 — at that point the cube is rising 12 units per unit of x. Its curvature is 6·2 = 12 — the slope itself is growing at 12 per unit. Three functions, one line each, no .backward() anywhere. PyTorch can produce these too, but in JAX it is the foundation rather than a bolt-on.
A second, gentler example — the one the chapter’s lab opens with — is f(x) = x² + 3x at x = 2:
Under the hood. JAX implements grad with reverse-mode automatic differentiation — the same chain rule that makes backpropagation work in Lesson 03, re-packaged as a function transform. When grad wraps f, it runs f over traced values and composes one local derivative rule per primitive operation, building the backward pass for you. The practical differences from .backward(): nothing is stored on tensors, the graph is not kept alive after the call, and differentiating again is just another wrap.
It composes, which is the point. jax.grad(jax.grad(f)) is a second derivative; add one more wrap for the third. jax.value_and_grad(loss_fn) returns the loss value and the gradient tree in a single pass, which is exactly what a training step needs. argnums chooses which arguments to differentiate — grad(loss, argnums=(0, 1)) hands back a tuple of gradients, one per argument. And for vector-valued outputs, jax.jacrev and jax.jacfwd build full Jacobians the same way.
The constraint is the same purity every other verb demands: no printing (it runs at trace time), no mutation of outside state, and randomness only through explicit keys. And there is one more rule worth internalizing immediately.
Pick a function and slide the point. The orange line is the exact derivative — the same value jax.grad returns. The dashed green line is the finite difference that sleeps in every calculus textbook; shrink h and watch it converge to the tangent.
On a parabola the central difference is exact at any h — its error term carries f‴ = 0. Switch to x³ and the gap returns: central difference error is O(h²), so h = 0.5 leaves a 0.25 error while h = 0.01 leaves 0.0001. Autodiff never has that gap.
Start from the definition and let the algebra do the work. For f(x) = x² + 3x at x = 2:
The cancellation is exact, which is why autodiff reports 7.0 while a numerical estimate carries an error. Plug in a real h to see it:
Two lessons hide in these decimals. First, the finite difference is a check, not the algorithm — shrink h and it converges to the analytic answer, but autodiff never needs it. Second, the error size tells you the flavor of the difference: forward differences are O(h), central differences are O(h²), and the lab’s gap readout shows both live.
jit does not make your function faster — it stops running your function at all. On the first call JAX traces it into a graph and XLA compiles that graph to machine code. Later calls with the same shapes execute the compiled artifact directly, with no Python in the loop.
Call one: the trace. JAX runs your function with abstract values — shape and dtype only, no real numbers. It records every primitive operation into a functional intermediate representation (a jaxpr), and hands that graph to XLA (Accelerated Linear Algebra), Google’s compiler for GPUs and TPUs. XLA fuses elementwise chains, eliminates redundant memory copies, and emits optimized machine code. The result is cached under a key made of the argument shapes and dtypes.
Calls two through ten thousand: the cache. Same shapes means the compiled executable still applies, so JAX skips tracing, skips compilation, and skips Python. The cost structure flips: pay once at compile time, then near-zero dispatch on every call. Whether that trade wins depends on how many times you call:
When to reach for it, when to run away. JIT pays for training steps, inference, and any function called more than once on similar-shaped inputs. It hurts when a computation runs once (compilation costs more than the work), during debugging (tracing hides the real execution), and when control flow depends on traced values.
That last restriction is the one that changes how you write code. If the branch condition is a traced array, Python’s if cannot decide — the value does not exist yet, only its shape. The functional replacements are:
None of these are optional decorations. They are the price of letting the compiler see the entire loop as one program — and they are why a JAX training step can run 10,000 times without paying Python overhead on any of them.
Three field notes from the source’s deployment checklist. Warm up before you benchmark — the first call includes compilation and will poison any average that includes it. Avoid Python loops over JAX arrays inside a jitted function; each iteration re-enters the traced graph, and the fix is jax.lax.scan / fori_loop (or vmap, next chapter). And JAX pre-allocates about 75% of GPU memory by default, which surprises anyone sharing a device — set XLA_PYTHON_CLIENT_PREALLOCATE=false to disable it.
Press the call buttons and watch the counters. The Python body runs only when a new shape forces a retrace — every other call is a cache hit. Then break it with the Python if.
JAX runs the function with abstract values — shapes and dtypes only. It records the ops as a jaxpr. Python `print()` and side effects happen here, once.
The traced graph goes to XLA, which fuses elementwise chains, removes copies and emits machine code. This is the ~1.2 s slow part.
The compiled executable runs on the accelerator. Later calls with the same shapes hit the cache: Python is skipped entirely.
Changing the batch size from 32 to 64 is not a bug — it is a new shape, so JAX retraces and recompiles. Real training loops keep every batch the same shape (drop or pad the last one) exactly so the cache stays warm.
A @jax.jit function contains print('step'). You call it 100 times with the same input shapes. How many times does 'step' appear in the output?
Forget the batch dimension. Write the function for a single example, then let vmap lift it over an axis — no loop, no reshaping, no batch-dimension threading. And because vmap composes with grad and jit, the batch version is one fused kernel.
The source’s example is as small as it gets — a dot product plus a bias, written for one input vector:
in_axes=(None, 0) is the whole contract, read position-by-position against the arguments. None for params means do not map — the weights are shared across the batch. 0 for x means map over axis 0 — each row is a separate call. out_axes defaults to 0, so the results come back stacked in the same order. Need to batch several arguments at once and see the semantics written out? in_axes=(None, 0, 0) shares the first argument and maps the rest — exactly the pattern for (params, x, y).
Worked numbers. Take w = [0.2, −0.4, 0.6], b = 0.5, and batch three rows. The single-example function does one dot product each; the readout below computes them by hand:
This is not syntactic sugar. A Python loop over the three rows would dispatch each dot product separately; vmap lowers the mapped function to vectorized primitives, so the same computation runs as one batched kernel — the source quotes 10–100× over a naive Python loop, and because the batch dimension is just another traced axis, it fuses with everything else.
The composition is where vmap earns its keep:
Two practical notes before the lab. If the mapped function uses randomness, vmap cannot invent independent draws — the key itself must be batched, one distinct key per example (Chapter 06 shows exactly how to split them). And for multiple devices, jax.pmap(f, axis_name="devices") is the same idea one level up: replicate the single-device function across every accelerator, split the batch, and average gradients with jax.lax.pmean. That is the programming model behind Gemini-scale runs: write the one-device version, wrap it, and let the framework place it on thousands of TPU v5e chips.
One pure function, seven wrappers. Pick a composition and read what it changes about the signature, where it runs, and the numbers it produces on the fixed example w = [0.2, −0.4, 0.6], b = 0.5.
Per-example gradients in one line. Their mean is the batch gradient: a free consistency check, shown live below.
The order of wrappers matters for shapes and cost, not for the math: jit(vmap(f)) and vmap(jit(f))compute the same numbers. Try the vmap∘grad preset last — the “mean of per-example gradients” line is the identity the lesson’s exercise builds on.
predict handles a single x of shape (3,) and returns a scalar. batch_predict = jax.vmap(predict, in_axes=(None, 0)). You call batch_predict(params, X) with X of shape (128, 3). What shape comes back?
Randomness is state, and JAX does not allow hidden state. So every random call names its source: a key you pass in, split, and never reuse. The rule feels like paperwork for a week; then you debug a multi-GPU run and realize it is the reason results reproduce.
In NumPy or PyTorch you seed a global generator once and every later call consumes the stream in order. JAX has no such stream. The key is the randomness:
Modern JAX prefers the typed constructor key = jax.random.key(42) — same behavior, but the value carries a distinct key<fry> dtype so it cannot be mistaken for an array of data. The legacy PRNGKey used throughout the source still works everywhere. Under the hood both are threefry, a counter-based pseudorandom function: a key selects a stream, and a draw walks a counter along it.
To get more randomness, split. random.split(key) deterministically produces two child keys that are unrelated to each other and to the parent:
Why you must never reuse a key. This is the part that bites everyone, so the lesson proves it with the real numbers. Because a draw always starts at position 0 of its key’s stream, the same key and shape replay the identical values — and the same key with a different shape replays an overlapping prefix:
So a reused key in initialization is not “a bit less random” — it is two weight matrices built from the same numbers, which silently undoes the symmetry breaking from Lesson 08. A reused key for dropout means the same mask every step. The discipline that fixes it is small: split once per consumer, and re-split once per step. Parameters get keys at init; the training loop keeps one running key and splits a fresh subkey for every epoch’s shuffle; each stochastic layer splits its own key on the way down. Keys flow through the program exactly like parameters do:
A JAX key is not a seed you sprinkle once — it is a value you spend. Bottom row: one split per draw, every leaf a different key. Flip the toggle to see what reuse actually produces.
Watch distinct samples as you deepen the tree: it equals the draw count when each consumer holds its own key. Now toggle reuse — the count collapses to one.
A key is just two uint32 words. Run k1, k2 = random.split(random.PRNGKey(42)) in a notebook and print them; here is what comes out:
The tree in the lab is drawn from exactly this table. Notice what makes it reproducible: nothing here reads a clock, a thread id, or global process state. The seed fixes the root, and the program’s data flow fixes everything else — so the same code produces the same weights on one GPU, on 64 TPUs, and in a re-compiled binary. A global seed cannot promise that once several workers consume the stream in nondeterministic order.
Two identical models are initialized with the same key: params_a = init(PRNGKey(0)) and params_b = init(PRNGKey(0)). Every draw inside init uses that key. What comes out?
Everything in this lesson composes into one step: compute the loss and its gradient in a single pass, let Optax transform the gradients, apply the update, and return the new state. Then compile the whole thing and let it run ten thousand times without Python.
The source builds a 3-layer MLP for MNIST: 784 inputs, two hidden layers of 256 and 128, 10 output classes. No class, no module — just a function that returns a pytree of parameters, and pure functions for the forward pass and the loss.
def init_params(key):
k1, k2, k3 = random.split(key, 3)
scale1 = jnp.sqrt(2.0 / 784)
scale2 = jnp.sqrt(2.0 / 256)
scale3 = jnp.sqrt(2.0 / 128)
return {
'layer1': {'w': scale1 * random.normal(k1, (784, 256)),
'b': jnp.zeros(256)},
'layer2': {'w': scale2 * random.normal(k2, (256, 128)),
'b': jnp.zeros(128)},
'layer3': {'w': scale3 * random.normal(k3, (128, 10)),
'b': jnp.zeros(10)},
}
def forward(params, x):
x = jnp.dot(x, params['layer1']['w']) + params['layer1']['b']
x = jax.nn.relu(x)
x = jnp.dot(x, params['layer2']['w']) + params['layer2']['b']
x = jax.nn.relu(x)
x = jnp.dot(x, params['layer3']['w']) + params['layer3']['b']
return x
def loss_fn(params, x, y):
logits = forward(params, x)
one_hot = jax.nn.one_hot(y, 10)
return -jnp.mean(
jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1)
)Now the step itself. This is the function the whole lesson has been building toward, and it is nine lines:
optimizer = optax.adam(learning_rate=1e-3)
@jax.jit
def train_step(params, opt_state, x, y):
loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
return params, opt_state, loss
@jax.jit
def accuracy(params, x, y):
logits = forward(params, x)
preds = jnp.argmax(logits, axis=-1)
return jnp.mean(preds == y)Notice what is missing: no .zero_grad(), no .backward(), no .step(). Those three calls existed to manage hidden state — clearing the old gradients, writing the new ones into tensors, mutating the weights. With state explicit, the entire update is one composed function call: gradients in, new params out. The optimizer state is a pytree too, which is why it is threaded through and returned alongside the parameters.
Strip the network down to one weight and one bias and the whole step fits on a line you can check by hand: ŷ = w·x + b, L = (ŷ − y)², with x = 2, y = 5, starting at w = 0.5, b = 0, and lr = 0.05. The gradients are ∂L/∂w = 2(ŷ − y)·x and ∂L/∂b = 2(ŷ − y):
This is the same nine-line train_step as the MNIST version, with the network replaced by one multiply. The shape of the computation never changes — and neither does the fact that every value is returned, not written.
The loop around it is ordinary Python: shuffle the data with a fresh subkey each epoch, slice fixed-shape batches, call the compiled step. The only JAX-specific lines are the key split and the fact that the step’s return value is reassigned:
key = random.PRNGKey(0)
params = init_params(key)
opt_state = optimizer.init(params)
batch_size, n_epochs = 128, 10
for epoch in range(n_epochs):
key, subkey = random.split(key) # fresh shuffle key
perm = random.permutation(subkey, len(X_train))
X_shuffled, y_shuffled = X_train[perm], y_train[perm]
epoch_loss = 0.0
n_batches = len(X_train) // batch_size
for i in range(n_batches):
start = i * batch_size
xb = X_shuffled[start:start + batch_size]
yb = y_shuffled[start:start + batch_size]
params, opt_state, loss = train_step(params, opt_state, xb, yb)
epoch_loss += loss
train_acc = accuracy(params, X_train[:5000], y_train[:5000])
test_acc = accuracy(params, X_test, y_test)
print(f"Epoch {epoch + 1:2d} | Loss: {epoch_loss / n_batches:.4f} | "
f"Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")One linear-regression training step on x = 2, y = 5 with w = 0.5, b = 0 and lr = 0.05. Every press returns a new pytree; the old one is never touched.
The button runs jax.value_and_grad(loss_fn)(params), then applies the update and returns the result. Nothing is mutated — this is a function from state to state.
| step | w | b | ŷ | loss |
|---|---|---|---|---|
| — | 0.5000 | 0.0000 | 1.0000 | 16.0000 |
Watch the loss column: 16 → 4 → 1 → 0.25 → 0.0625. Each step quarters the error because the learning rate 0.05 is exactly tuned to this quadratic — the same behavior the chapter proves by hand.
The ecosystem around the primitives. JAX gives you grad, jit, vmap and pytrees. Everything else is a library, and the standard ones are worth knowing by name:
When to use JAX, honestly. The source keeps a table instead of a slogan, and it is worth reading as written:
Answer before you look. The jit question and the key-reuse question are the two that separate a list of magic names from a mechanism you can debug.
0 / 5 answered · 0 correct
01What is the fundamental design difference between PyTorch and JAX?
02What does jax.jit do?
03What does jax.vmap do?
04How does JAX handle model state (weights) differently from PyTorch?
05When would you choose JAX over PyTorch?
Click a card to swap the lazy description for what it actually means.
Four problems with exact numbers and correct-looking code — add dropout with explicit keys, compute per-example gradients with vmap, write a depth-agnostic forward pass, and measure what jit actually buys you. Try first; a worked answer is one click away.
The functional shape is a forward function that takes the key as an argument and returns the updated key alongside the output — state in, state out, exactly like parameters. Per layer: key, subkey = random.split(key); mask = random.bernoulli(subkey, keep, x.shape); x = jnp.where(mask, x, 0.0) / keep. The / keep is inverted dropout: it rescales so the expected activation is unchanged, and it is why nothing needs to be corrected at eval time. Evaluation is not a flag mutation (no model.eval()): in JAX you write a second pure function without the mask, or make the training flag static — jax.jit(forward, static_argnames=('train',)) — because a Python if on a traced value cannot decide, while a Python if on a static flag is resolved before tracing. Comparison: on MNIST the test accuracy typically lands within ±0.5% of the ~97% no-dropout baseline (dropout often does not help on this small, clean dataset); the exercise's real payoff is the key plumbing, which is the source's stated point. Numeric check on key independence: with keep = 0.8, a 256-unit mask zeroes ~51 units (Binomial(256, 0.2), mean 51.2, sd 6.4). Two independent masks agree at each unit with probability 0.8² + 0.2² = 0.68, so they agree everywhere with probability 0.68²⁵⁶ ≈ 10⁻⁴³. Reuse the key — or fail to reassign key = subkey in the training loop — and that exponent becomes 1: the same mask, every step, forever.
Write a per-example loss, then map grad over the batch: def loss_one(params, x, y): logits = forward(params, x); return -jnp.sum(jax.nn.log_softmax(logits) * jax.nn.one_hot(y, 10)); per_example = jax.vmap(jax.grad(loss_one), in_axes=(None, 0, 0))(params, xb, yb); norms = jax.vmap(lambda g: jnp.sqrt(sum(jnp.sum(leaf ** 2) for leaf in jax.tree.leaves(g))))(per_example). The largest norms belong to examples the model gets wrong most — high loss means a large gradient — weighted by the example's input scale. Check the mechanism on a linear model, where ŷ = w·x and L = (ŷ − y)², so ∇_w L = 2(ŷ − y)·x and its norm is 2·|error|·‖x‖. With w = [0, 0], b = 0: example 1 has x = [1, 0], y = 1 → error −1 → ∇ = (−2, 0), norm 2. Example 2 has x = [10, 0], y = −1 → error +1 → ∇ = (20, 0), norm 20. Same error, 10× the input norm, 10× the gradient norm. On MNIST the top-norm examples are typically confidently misclassified digits with a lot of inked pixels — the ones a regularizer or a clipping rule most wants to see. That is exactly why DP-SGD clips per-example gradients before averaging: an example that shouts cannot dominate the update. Sanity check: for this loss, mean over examples of the per-example gradients equals the gradient of the mean loss — the lab's vmap∘grad preset computes both sides live and they match to float32.
def mlp_forward(params, x): leaves = jax.tree.leaves(params); depth = len(leaves) // 2; for layer in range(depth): block = params[f'layer{layer + 1}']; x = jnp.dot(x, block['w']) + block['b']; if layer < depth - 1: x = jax.nn.relu(x); return x. Each layer contributes one weight and one bias, so leaves/2 is the depth; for the source's tree the leaves come back in sorted key order — [l1.b, l1.w, l2.b, l2.w, l3.b, l3.w] — because JAX sorts dict keys lexicographically — 6 leaves → depth 3. The Python if is safe under jit because depth is a Python integer derived from the pytree structure, not a traced value. Two honest notes. First, the sorted-key traversal is a real footgun: a model with ten layers named layer1…layer10 would visit layer10 before layer2; pad the names (layer01) or store the layers in a list, where order is explicit. Second, the generic version hides a real assumption: every block has exactly a 'w' and a 'b' of matching shapes. If you want it to be robust, traverse blocks = leaves[0::2] and derive shapes from each block instead of assuming them — the source's exercise is about the pattern, not about bulletproofing. Parameter-count check with the source's tree: 200,960 + 32,896 + 1,290 = 235,146 across three layers, and mlp_forward must produce identical logits to the manual version for every input.
Methodology first, because three details make naive timing wrong. (1) JAX is asynchronous: a call returns after dispatching, not after the kernels finish, so time a block of calls and end with jax.block_until_ready(out) — or call out.block_until_ready() — before stopping the clock. (2) Warm up the jitted function once before timing, so you measure steady state, and time the very first call separately to capture the trace+compile overhead. (3) Keep shapes identical across timed calls; one changed batch size silently adds a compile to your measurement. Pattern: def timeit(fn, n=100): fn(); jax.block_until_ready(fn()); start = time.perf_counter(); out = None; for _ in range(n): out = fn(); jax.block_until_ready(out); return (time.perf_counter() - start) / n. Representative numbers — a worked example, not a measurement — are 8.0 ms eager vs 0.4 ms compiled per step and 1.2 s of compilation overhead: a 20× steady-state speedup that breaks even after ~158 steps and saves ~75 s over 10,000 steps (~15× end to end). Your numbers will differ: on CPU the eager Python overhead dominates and the ratio is often double-digit; on GPU the eager version overlaps much of that overhead with compute, so the steady-state ratio can be smaller — but the JAX win grows with step complexity and is essential on TPUs, where eager dispatch per op is far too slow. Report all three numbers honestly: first-call time, steady-state times, and the break-even step count.
You do not need to master these here. Each one gets a proper treatment in its own lesson; the one-line meaning is enough to keep reading. Orange dotted underlines in the prose point back to this list.
This lesson is a port of an open course. Everything here traces back to it — and the next step is running the code yourself.
Lesson text adapted from AI Engineering from Scratch (Phase 03, Lesson 12) and the Math Foundations Notebook reference build. The seven labs (the derivative machine canvas, the key-splitting tree canvas with real JAX 0.6.2 threefry values, the trace-timeline stepper, the params-in/params-out pure-state stepper, the transformation board, the PyTorch-to-JAX translation board, and the immutability console) are original to this page, as are the eager-dispatch overhead arithmetic, the x² + 3x limit derivation with forward and central finite-difference checks (7.001 versus 7.000000 and the cube's 12.000001), the worked pytree update, the jit amortization table with its 158-step break-even, the worked vmap dot products and the per-example-gradient identity, the real PRNG split tables and same-key stream replays, the linear-regression training step and its 0.5 error contraction, the softmax cross-entropy hand check, the dropout mask-collision arithmetic, and the memory hooks. Every number shown is computed live by the labs or verified by hand in the prose.