jacfwd vs jacrev

Put each output's derivatives in a row

A Jacobian collects the first derivatives of a function with several outputs. If the input is a vector of length d and the output is a vector of length m, its shape is (m, d). Entry (i, j) tells you how output i changes when input j changes.

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])
J_fwd = jax.jacfwd(f)(x)
J_rev = jax.jacrev(f)(x)
# Both give [[1.0, 1.0],
#            [3.0, 2.0]]

The first output is a sum, so both derivatives in its row are 1. For the product, changing x[0] has derivative x[1]=3, and changing x[1] has derivative x[0]=2.

Two ways to build the matrix

Forward-mode jax.jacfwd builds the Jacobian from input directions: one coordinate direction selects one column. Reverse-mode jax.jacrev builds it from output weights: selecting one output gives one row. JAX batches these derivative calculations, so this is not a promise of one separate Python call or device operation per row or column.

For a full Jacobian, forward-mode is often a good choice when there are few inputs and many outputs. Reverse-mode often suits many inputs and few outputs. These are guidelines, not timing guarantees; the function, compiler and hardware also affect cost. Both compute the same mathematical Jacobian for the smooth function here, though floating-point results need not be bitwise identical.

For array inputs and outputs, output axes come before input axes. A scalar output contributes no axes: for vector input x, both methods return shape (d,), just like jax.grad. There is no extra length-one output axis. The exercise below has three outputs, so it needs three rows.

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.