grad vs vjp

Recover a gradient from a pullback

jax.grad gives the gradient of a scalar loss. jax.vjp gives you the loss value and a function called a pullback. For a real scalar loss, calling that pullback with a weight of one gives the same mathematical gradient. A weight of three gives three times the gradient.

import jax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0])

def loss(w):
    return jnp.sum((w * x) ** 2)

value, pullback = jax.vjp(loss, 1.0)
# value is 5.0
pullback(1.0)[0]       # 10.0
jax.grad(loss)(1.0)    # 10.0
pullback(3.0)[0]       # 30.0

Here loss(w) = 5*w**2, whose derivative at w=1 is 10. The weight passed to the pullback is called a cotangent. It scales that derivative; it does not change the w at which the derivative is evaluated.

Match the output, then unpack the inputs

The loss output has shape (), so the cotangent must also be scalar. A Python float works in this example. A zero-dimensional JAX array also works: jnp.ones_like(value) supplies one with the output's shape and dtype. A length-one array has shape (1,) and is not interchangeable.

The pullback returns a tuple with one entry for each argument passed to jax.vjp. Our inner function accepts only w, so [0] extracts its scalar derivative. Capturing b and x from the surrounding function keeps them fixed for this derivative.

Compute derivatives by following the operations

Automatic differentiation computes derivatives by applying the chain rule to a program's operations. It does not estimate them by evaluating nearby inputs, as finite differences do, or rearrange the whole expression using symbolic algebra.

For f(g(x)), the chain rule combines the derivative of g with the derivative of f at g(x). Repeating this process lets a framework differentiate a long sequence of operations.

Forward mode follows an input change

For a function with n inputs and m outputs, the Jacobian J is an (m, n) matrix of derivatives. Entry (i, j) describes how output i changes with input j.

Forward-mode differentiation computes Jv: the output's rate of change when the input moves in direction v. For example, if f(x) = sum(x_i^2), then at x = [1, 2] the gradient is [2, 4]. In direction v = [3, 0], the rate of change is 2 * 3 + 4 * 0 = 6.

This is a directional derivative. It gives one Jacobian column only when v selects one input coordinate, such as [1, 0].

Reverse mode works backward from an output

Reverse mode computes J^T u, where u weights the outputs. For a scalar loss, choosing u = 1 gives the gradient with respect to all inputs. This makes reverse mode useful for training: a model can have many parameters but just one loss.

Forward mode can be useful when you need sensitivity to only a few input directions. Neither mode needs to construct the full Jacobian for these products. That can save substantial memory. Higher-order differentiation can similarly compute a Hessian-vector product without storing the full matrix of second derivatives.

Changing gradient flow changes what is differentiated

A stop-gradient operation keeps a value in the forward calculation but treats it as constant during differentiation. A custom derivative supplies a derivative rule explicitly. It should match the intended function unless the exercise asks for a surrogate, such as an approximation used during training.