vjp Basics

Send an output weight backward

For a scalar loss, multiplying the loss by a constant also multiplies its gradient by that constant. jax.vjp lets you supply that weight when running the derivative calculation backward. The weight is called a cotangent.

import jax
import jax.numpy as jnp

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

x = jnp.array([1.0, 2.0, 3.0])
value, pullback = jax.vjp(loss, x)
# value is 14.0
pullback(1.0)[0]   # [2.0, 4.0, 6.0]
pullback(-2.0)[0]  # [-4.0, -8.0, -12.0]

value is the ordinary loss. pullback is a function that maps a weight on the output to weights on the inputs, using derivatives at this x. Here the loss derivative is 2*x, so each call scales that vector. The pullback returns a tuple with one entry per input argument. Since loss has one argument, [0] extracts its only entry.

What makes it a vector-Jacobian product?

For a function with several outputs, supply one weight per output. If J[i, j] is the derivative of output i with respect to input j, the result's entry j is sum_i cotangent[i] * J[i, j]. This is the derivative of a weighted sum of outputs, with those weights held fixed. JAX can compute this product without constructing the whole Jacobian matrix.

For example, if f(x) = [x[0]**2, 3*x[1]], then at x=[2, 4] the output weights [2, -1] produce input weights [8, -3].

For array outputs, the cotangent must have the output's shape and a compatible dtype. A scalar output has shape (): both 1.0 and jnp.array(1.0) work here, but jnp.array([1.0]) has the wrong shape. You can reuse a pullback with different cotangents at the same input. To evaluate derivatives at a new input, call jax.vjp again.

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.