Autodiff

Call jax.grad on a function and JAX does not perturb the inputs and divide by a small number. It looks at the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, the recorded program from chapter 2, and derives a new program from it: one that computes the exact gradient by the chain rule, term for term. There is no step size to tune and no truncation error to worry about, because nothing was ever approximated. The rule for differentiating each primitive (a dot_general, a sin, a reduce_sum) is registered once, and grad composes those rules automatically along the trace.

The usual finite-difference intuition, nudge x by epsilon and measure how y moved, has nothing to do with what happens here. jax.grad(f) expects f to return a scalar: a loss, not a vector of predictions. Ask for the gradient of something vector-valued and JAX complains, because the gradient of a vector output isn't a vector at all, it's a Jacobian, and later chapters build those directly with other tools. When you also need the loss value itself, which you almost always do for logging, jax.value_and_grad gets you both from a single forward-and-backward pass instead of running the function twice.

Because grad is a transformation from one program to another, it composes with every other transformation the same way. grad(grad(f)) differentiates twice. Wrap a vmapped function in grad, or a scanned one, and the derivative flows through correctly, not because JAX special-cased those combinations, but because a transformed program is still just a program, and grad doesn't need to know how it got built.

grad is a program transformation, not a numerical method.

import jax
import jax.numpy as jnp

def loss(w, x, y):
    return jnp.mean((x @ w - y) ** 2)

g = jax.grad(loss)(jnp.ones(3), jnp.eye(3), jnp.zeros(3))
print(g)   # exact, to machine precision: [0.6667 0.6667 0.6667]

Underneath grad sit two more primitive transformations, and picking between them by hand is worth understanding even if you rarely call them directly. jax.jvp pushes a tangent vector forward through the program: for every input direction you care about, it costs one extra forward-like pass. That's cheap when you have a few inputs and many outputs, a Jacobian-vector product for a small parameter feeding a large output, and expensive the moment your input space gets big.

jax.vjp runs the opposite direction. It does one forward pass, but it also has to remember enough of that pass, its residuals, to walk backward and pull a cotangent through every operation. The backward pass costs about as much as the forward one, and critically, that total cost does not grow with the number of inputs. This is why deep learning lives in reverse mode almost exclusively: a model can carry a billion parameters and one scalar loss, and reverse mode differentiates all billion of them in roughly two forward passes' worth of work.

jacfwd and jacrev, the two ways to build a full Jacobian, are not separate machinery. Each is vmap wrapped around jvp or vjp respectively, batching the single-direction cost across every direction at once. Which one to reach for is a question about the Jacobian's shape: forward mode when there are few inputs and many outputs, reverse mode when it's the other way around.

The bill reverse mode runs up is memory, not compute: whatever the backward pass needs to remember from the forward pass has to live somewhere until it gets used. The last section of this chapter comes back to that bill directly.

Every rule grad applies by default comes from JAX's built-in table of how to differentiate primitives, and most of the time that table is exactly what you want. Sometimes it isn't. jax.custom_vjp lets you throw out the automatic backward rule for a function and supply your own: because the automatic one is numerically unstable, because you know an algebraically cheaper form, or because the function hides work, a kernel, a call to something the tracer can't see through, that autodiff has no rule for at all. The kernel path's flash-attention backward (at /s/kernels) is built exactly this way, replacing a rule the tracer could never have derived on its own.

Take log1p(exp(x)), a softplus that looks numerically stable on its face. Its automatic gradient is fine for small x, but at large x, exp(x) overflows toward infinity while the chain rule needs to multiply that infinity by a value collapsing toward zero, and the result is nan. The function's value comes out correctly; the gradient the automatic rule derives from it does not. Nothing is broken in JAX here: the algebra it's following really does produce that indeterminate form. The fix is to hand it a gradient rule you derived by hand, one that never routes through the unstable intermediate at all.

import jax
import jax.numpy as jnp

@jax.custom_vjp
def log1pexp(x):
    return jnp.log1p(jnp.exp(x))

def fwd(x):
    return log1pexp(x), jax.nn.sigmoid(x)   # residual: the stable gradient

def bwd(sig, ct):
    return (ct * sig,)

log1pexp.defvjp(fwd, bwd)
print(jax.grad(log1pexp)(100.0))   # 1.0; the automatic rule returns nan here

Reverse mode's residuals are usually intermediate activations: everything the backward pass needs to reconstruct the chain rule. For a deep network that's one saved tensor per layer, and it adds up fast. jax.checkpoint, JAX's name for gradient checkpointing (remat, if you've met the term elsewhere), makes a different trade: instead of storing a block's activations, it discards them and recomputes that block during the backward pass. You spend more FLOPs to spend fewer bytes. It's the identical trade flash attention makes at the kernel level, applied here at the level of whole layers instead of attention blocks.

jax.lax.stop_gradient does something unrelated: it cuts the tape at a specific point, so no gradient flows through that value afterward, no matter what happens to it downstream. Target networks in reinforcement learning, straight-through estimators, freezing a pretrained encoder while still using its output: all three are the same operation underneath, telling autodiff that a particular value counts as a constant for the backward pass, even though it came from real computation.

import jax
import jax.numpy as jnp

W = jnp.ones((256, 256)) * 0.01

def block(x):
    for _ in range(8):
        x = jnp.tanh(x @ W)
    return x.sum()

g_full = jax.grad(block)(jnp.ones(256))
g_remat = jax.grad(jax.checkpoint(block))(jnp.ones(256))
print(jnp.allclose(g_full, g_remat))   # True; the remat pass stored no intermediates