jax.linearize Primitive

Reuse a derivative at one point

Suppose you want to see how a function responds to several input directions, all starting at the same x. jax.linearize(f, x) returns two things: the value f(x) and a function that applies the derivative at that x to a direction.

import jax
import jax.numpy as jnp

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

x = jnp.array([1.0, 2.0])
value, linear = jax.linearize(f, x)
# value = 5.0
linear(jnp.array([1.0, 0.0]))   # 2.0
linear(jnp.array([0.0, 1.0]))   # 4.0
linear(jnp.array([-1.0, 2.0]))  # 6.0

The gradient at x=[1, 2] is [2, 4]. Each call takes its dot product with the supplied direction. More generally, the returned function computes J(x) @ v, a Jacobian-vector product. Its output has the same shape as the function's output.

A direction is not a new evaluation point

The returned function is linear in v, but f itself need not be linear. For small epsilon, the first-order estimate is f(x + epsilon*v) ≈ value + epsilon*linear(v). The derivative function does not return f(x + v) or include the original value automatically.

For example, with x=[1, 2] and v=[1, 0], the estimate for f(x+v) is 5 + 2 = 7, while its actual value is 8. The difference comes from the quadratic term. To compute derivatives at a different input point, call jax.linearize again at that point.

Reuse has a memory cost

jax.jvp(f, (x,), (v,)) returns the value and one tangent immediately. jax.linearize instead returns a callable for later tangents. It keeps intermediate information needed by that callable, rather than building and caching a full Jacobian matrix. Storage can grow with the original computation, and each application still performs derivative operations.

This can be useful when directions arrive one at a time. If you already have all directions, batching JVPs with jax.vmap is another option that can avoid linearize's stored-intermediate cost. Choose based on the workload rather than assuming every reused call is cheap.

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.