Jacobians, JVPs and VJPs

Physics, optimization and research code need Jacobians, Hessian-vector products and custom gradient flow.

grad is a special case of two primitives:

Forward mode, jax.jvp(f, (x,), (v,)), pushes a tangent v through f and returns (f(x), J v). Cost is about one evaluation per input direction.
Reverse mode, jax.vjp(f, x), returns f(x) and a function computing u^T J for any cotangent u. Cost is about one evaluation per output direction. grad is a VJP with u = 1.

From these JAX builds full Jacobians:

jax.jacfwd(f) uses forward mode, best for "tall" Jacobians (few inputs, many outputs).
jax.jacrev(f) uses reverse mode, best for "wide" ones (many inputs, few outputs).
jax.hessian(f) is jacfwd(jacrev(f)), the efficient forward-over-reverse composition.

Materializing a Hessian is quadratic in parameter count, so for large models compute Hessian-vector products instead: jax.jvp(jax.grad(f), (x,), (v,))[1]. This costs only a few gradient evaluations.

jax.lax.stop_gradient(x) treats x as a constant for differentiation, useful for target networks, straight-through estimators and EM-style updates.

All of these compose with jit and vmap, e.g. vmap(jacrev(f)) for batched Jacobians. Mixing modes is how JAX makes higher-order derivatives practical.

Example

import jax
import jax.numpy as jnp

def f(x):                       # R^2 -> R^3
    return jnp.array([x[0] * x[1], jnp.sin(x[0]), x[1] ** 2])

x = jnp.array([1.0, 2.0])
print(jax.jacfwd(f)(x).shape)   # (3, 2)
print(jnp.allclose(jax.jacfwd(f)(x), jax.jacrev(f)(x)))  # True

y, jv = jax.jvp(f, (x,), (jnp.array([1.0, 0.0]),))
print(jv)                       # first column of J: [2. 0.5403023 0.]

def g(x):
    return jnp.sum(x ** 3)
print(jax.hessian(g)(x))        # [[ 6.  0.] [ 0. 12.]]

v = jnp.array([1.0, 1.0])
hvp = jax.jvp(jax.grad(g), (x,), (v,))[1]
print(hvp)                      # [ 6. 12.] without building the Hessian

h = lambda x: x * jax.lax.stop_gradient(x)
print(jax.grad(h)(3.0))         # 3.0 (not 6.0)

Remember

JVPs (forward) cost per input direction; VJPs (reverse) per output direction
jacfwd for tall Jacobians, jacrev for wide ones, hessian = forward-over-reverse
Use Hessian-vector products instead of full Hessians for large models

Watch out for

Building a full Hessian for millions of parameters instead of using Hessian-vector products.

Check yourself

f maps 1,000,000 parameters to a scalar loss. Which is the efficient way to get its full gradient?

Official docs: docs.jax.dev/en/latest/301/cookbook.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks