Higher-Order Gradients and Jacobians

Because jax.grad returns a regular function, you can differentiate it again — and again. This gives you higher-order derivatives with zero extra code.

Higher-Order Gradients

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x)

df = jax.grad(f)         # cos(x)
d2f = jax.grad(df)       # -sin(x)
d3f = jax.grad(d2f)      # -cos(x)
d4f = jax.grad(d3f)      # sin(x)

x = jnp.array(jnp.pi / 4)
print(f"f(x)   = {f(x):.4f}")     # 0.7071 (sin)
print(f"f'(x)  = {df(x):.4f}")    # 0.7071 (cos)
print(f"f''(x) = {d2f(x):.4f}")   # -0.7071 (-sin)
print(f"f'''(x)= {d3f(x):.4f}")   # -0.7071 (-cos)

Hessians: Second Derivatives of Multi-Variable Functions

import jax
import jax.numpy as jnp

def f(x):
    return x[0] ** 2 * x[1] + x[1] ** 3

# Hessian: matrix of second partial derivatives
hessian_fn = jax.hessian(f)
x = jnp.array([1.0, 2.0])
H = hessian_fn(x)
print(H)
# [[ 4.  2.]   d^2f/dx0^2 = 2*x1 = 4,  d^2f/dx0dx1 = 2*x0 = 2
#  [ 2. 12.]]  d^2f/dx1dx0 = 2*x0 = 2,  d^2f/dx1^2  = 6*x1 = 12

Jacobians: For Vector-Valued Functions

jax.grad only works on scalar-valued functions (it computes the gradient). For functions that output vectors or matrices, you need Jacobians:

import jax
import jax.numpy as jnp

def vector_fn(x):
    """Maps R^3 -> R^2."""
    return jnp.array([x[0] * x[1], x[1] ** 2 + x[2]])

x = jnp.array([1.0, 2.0, 3.0])

# jacrev: Reverse-mode Jacobian (efficient when output dim < input dim)
J_rev = jax.jacrev(vector_fn)(x)
print(J_rev)
# [[2. 1. 0.]    dy0/dx0=x1, dy0/dx1=x0, dy0/dx2=0
#  [0. 4. 1.]]   dy1/dx0=0,  dy1/dx1=2*x1, dy1/dx2=1
print(J_rev.shape)  # (2, 3)

# jacfwd: Forward-mode Jacobian (efficient when input dim < output dim)
J_fwd = jax.jacfwd(vector_fn)(x)
print(jnp.allclose(J_rev, J_fwd))  # True — same result, different algorithm

When to Use jacfwd vs jacrev

• jacrev (reverse mode): Efficient when the function has more inputs than outputs. Think "tall" Jacobian. This is what grad uses internally.
• jacfwd (forward mode): Efficient when the function has more outputs than inputs. Think "wide" Jacobian.