JVP with Two Input Directions

Change two inputs together

A Jacobian-vector product describes how a function responds along a chosen input direction. When the function has two arguments, provide a direction for each. JAX adds their contributions to the output tangent.

import jax
import jax.numpy as jnp

def f(a, b):
    return jnp.sum(a * b)

a = jnp.array([1.0, 2.0])
b = jnp.array([3.0, 4.0])
da = jnp.array([1.0, 0.0])
db = jnp.array([0.0, 1.0])
value, tangent = jax.jvp(f, (a, b), (da, db))
# value = 1*3 + 2*4 = 11
# tangent = sum(da*b + a*db) = 3 + 2 = 5

This differentiates the path f(a + epsilon*da, b + epsilon*db) at epsilon=0. The tangent is the rate of change along the whole path, not the exact difference between two evaluations a finite distance apart.

Match the input structure

The outer tuple holds one entry per positional argument. Matching lists also work. Each real floating-point input array needs a tangent with its shape and dtype. To hold an array fixed, supply jnp.zeros_like for its tangent; None is not a replacement for an array tangent.

The same rule extends to nested dictionaries and lists, called pytrees: every input array leaf needs a corresponding tangent leaf. This exercise uses two array arguments, so its structure is simply a pair of arrays. The separate diagram illustrates a more deeply nested tree. Its numbered leaves are placeholders, not the exercise's input values.

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.