jvp Basics

Follow a direction through a function

Suppose an input changes from x to x + epsilon*v, where epsilon is a small number. The vector v describes which input entries change and by how much relative to one another. A directional derivative tells you the output's rate of change along that direction.

jax.jvp returns the ordinary function value, called the primal, and this directional derivative, called the output tangent:

import jax
import jax.numpy as jnp

def square(x):
    return x ** 2

value, tangent = jax.jvp(square, (3.0,), (2.0,))
# value = 9.0
# tangent = 2 * 3 * 2 = 12.0

For small epsilon, square(3 + epsilon*2) is approximately 9 + epsilon*12. The tangent is a rate, not the value of the function at a displaced input. You do not need to choose an epsilon for JAX: it applies derivative rules to the operations in your function.

From tangents to a Jacobian-vector product

The Jacobian collects the derivatives of every output with respect to every input. Multiplying that matrix by v gives the output tangent, written J(x) @ v. JAX computes this product without needing to build the whole matrix. For a scalar loss, it is the dot product of the gradient and v; for an array output, it has that output's shape.

x = jnp.array([1.0, 2.0, 3.0])
v = jnp.array([1.0, 0.0, 0.0])
value, tangent = jax.jvp(lambda x: jnp.sum(x**2), (x,), (v,))
# value = 14.0; tangent = 2.0

Only the first input entry changes along this direction, so only its derivative contributes. A direction need not have length one: scaling v scales the tangent.

Pass the positional inputs and their tangents as matching tuples or lists. For one argument, (x,) and (v,) are convenient. Here the real floating-point arrays x and v must have the same shape and dtype.

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.