Automatic Differentiation

Hand-deriving gradients for a 100-million-parameter network is a non-starter. Every modern framework ships an automatic differentiation engine that:

Records each operation onto a computational graph.
Walks the graph in reverse to apply the chain rule.
Returns the gradient with respect to every input you asked about — typically the model parameters.

This chapter teaches the API; the rest of the book leans on it.

A worked example

We’ll differentiate

y = 2\,\mathbf{x}^\top \mathbf{x}

with respect to the column vector \mathbf{x}. The analytic gradient is \nabla_\mathbf{x} y = 4\mathbf{x} — a useful sanity-check target.

from jax import numpy as jnp

x = jnp.arange(4.0)
x

Array([0., 1., 2., 3.], dtype=float32)

Tracking gradients

We tell the framework to track operations on x and reserve a slot for its gradient:

Then run the forward pass — y is built from x, so the engine records the dependency:

y = lambda x: 2 * jnp.dot(x, x)
y(x)

Array(28., dtype=float32)

Backward pass

A single call walks the recorded graph backwards:

from jax import grad
# The `grad` transform returns a Python function that
# computes the gradient of the original function
x_grad = grad(y)(x)
x_grad

Array([ 0.,  4.,  8., 12.], dtype=float32)

The result lands in x.grad. Compare with the analytic answer, 4\mathbf{x}:

x_grad == 4 * x

Array([ True,  True,  True,  True], dtype=bool)

Resetting & re-using

Gradients accumulate by default — call .zero_() (or its equivalent) before computing a fresh gradient:

y = lambda x: x.sum()
grad(y)(x)

Array([1., 1., 1., 1.], dtype=float32)

For non-scalar y, the engine sums up gradients computed for each output element (or you supply weights):

y = lambda x: x * x
# grad is only defined for scalar output functions
grad(lambda x: y(x).sum())(x)

Array([0., 2., 4., 6.], dtype=float32)

Detaching from the graph

Sometimes we want a value treated as a constant in the backward pass — e.g., the auxiliary u below should not propagate gradients into x:

import jax

y = lambda x: x * x
# jax.lax primitives are Python wrappers around XLA operations
u = jax.lax.stop_gradient(y(x))
z = lambda x: u * x

grad(lambda x: z(x).sum())(x) == y(x)

Array([ True,  True,  True,  True], dtype=bool)

After detach() (or stop_gradient / lax.stop_gradient), the gradient flows around the detached tensor, not through it:

grad(lambda x: y(x).sum())(x) == 2 * x

Array([ True,  True,  True,  True], dtype=bool)

Gradients through control flow

Autograd doesn’t care about Python ifs and whiles — it records whichever ops actually executed. Here’s a function whose behavior depends on its input:

def f(a):
    b = a * 2
    while jnp.linalg.norm(b) < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

The number of while iterations and the branch taken both depend on the value of a.

…it just works

Run the function on a random scalar and ask for the gradient:

from jax import random
a = random.normal(random.PRNGKey(1), ())
d = f(a)
d_grad = grad(f)(a)

The gradient is correct even though the path through the function is data-dependent. Here f(a) ends up linear in a along whichever branch ran, so f'(a) = f(a) / a:

d_grad == d / a

Array(True, dtype=bool)

Recap

Mark inputs as needing gradients.
Run the forward pass — the engine records ops.
backward() (or grad()) walks the graph in reverse via the chain rule.
Gradients accumulate; reset between iterations.
detach / stop_gradient to break the graph.
Works through arbitrary Python control flow.