Build a Jacobian with VJPs

Select one output at a time

A vector-Jacobian product combines the derivatives of several outputs using weights called cotangents. To get the derivatives of just one output, give that output weight 1 and all the others weight 0. These weights form a standard basis vector.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.array([x[0] + x[1], x[0] * x[1]])

x = jnp.array([2.0, 3.0])
value, pullback = jax.vjp(f, x)
basis = jnp.eye(2, dtype=value.dtype)
row0 = pullback(basis[0])[0]  # [1.0, 1.0]
row1 = pullback(basis[1])[0]  # [3.0, 2.0]
J = jnp.stack([row0, row1])

basis[0] is [1, 0], so it selects the sum's derivatives. basis[1] is [0, 1], so it selects the product's derivatives. Stacking these results gives the Jacobian, with one row per output and one column per input.

The pullback returns a tuple with one entry per argument supplied to jax.vjp. Here there is only one input argument, so [0] extracts its array. Each cotangent must match the output's shape: length 2 in this example, and length 3 in the exercise below.

Reuse the pullback

Call jax.vjp once at the input where you want the Jacobian, then apply its pullback to each output basis vector. The three-row exercise can use a fixed Python loop. Under JIT, that fixed loop is traced and its operations enter the compiled computation.

You can also batch the basis vectors with jax.vmap(pullback)(basis)[0]. Here vmap batches output weights, not a batch of different input points. This is the row-building idea behind jax.jacrev. Batching can help performance, but it does not guarantee a speedup over a compiled loop or imply one separate device pass for every output.

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.