Chapter 4¶

Overview¶

Automatic differentiation is JAX's superpower. This chapter covers computing gradients with jax.grad.

-

What is Automatic Differentiation?¶

Automatic differentiation computes derivatives without:

Manual calculus
Finite differences approximation

Instead, JAX tracks operations and applies the chain rule automatically.

Why Not Manual Calculus?¶

# Manual
def f(x):
 return x**3 + 2*x**2 + 5*x + 1

# df/dx = 3x² + 4x + 5
def df_manual(x):
 return 3*x**2 + 4*x + 5

print(df_manual(2.0)) # 23.0

# But manually compute for complex functions? Impossible!
def complex_f(x):
 y = x**2
 z = jnp.sin(y) + jnp.exp(-y)
 return jnp.log(z + 2)

# df/dx by hand? Tedious and error-prone!

Why Not Finite Differences?¶

import jax.numpy as jnp

def f(x):
 return x**2

# Approximate derivative
eps = 1e-5
df_approx = (f(2.0 + eps) - f(2.0)) / eps
print(df_approx) # ~4.0 (should be exactly 4.0)

# Problems:
# 1. Numerical error (ε too small → rounding error)
# 2. Expensive (need f(x+ε) for each variable)
# 3. For 1M parameters

Automatic Differentiation is Better¶

import jax
import jax.numpy as jnp

def f(x):
 return x**2

# Exact derivative, automatic!
grad_f = jax.grad(f)
print(grad_f(2.0)) # 4.0 (exact)

# Works for any composition:
def complex_f(x):
 y = x**2
 z = jnp.sin(y) + jnp.exp(-y)
 return jnp.log(z + 2)

grad_complex = jax.grad(complex_f)
print(grad_complex(1.0)) # Exact derivative!

-

Forward-Mode vs Reverse-Mode¶

JAX supports both, but focuses on reverse-mode (backpropagation).

Reverse-Mode (Backpropagation)¶

import jax
import jax.numpy as jnp

def f(x):
 y = x**2 # y = x²
 z = jnp.sin(y) # z = sin(y)
 return z # return z

# Forward pass (normal computation):
# x=2 → y=4 → z=sin(4)

# Backward pass (gradient computation):
# dz/dz = 1
# dy/dx = 2x = 4
# dz/dx = dz/dy * dy/dx = cos(4) * 4

grad_f = jax.grad(f)
print(grad_f(2.0)) # cos(4) * 4 ≈ -2.614

Why reverse-mode? Efficient for many inputs, one output (like neural networks with single loss).

Forward-Mode¶

import jax

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

# Compute Jacobian-vector product (JVP)
# = "forward-mode differentiation"

y, vjp_fn = jax.linearize(f, 2.0)
# y = f(2.0) = 17
# vjp_fn = function to compute vjp

(Reverse-mode is primary; forward-mode used for specific cases)

-

Basic Usage: jax.grad¶

Simple Scalar Function¶

import jax
import jax.numpy as jnp

def f(x):
 return x**3 - 2*x + 5

# Create gradient function
grad_f = jax.grad(f)

# Compute gradient at x=2
gradient = grad_f(2.0)
# f'(x) = 3x² - 2
# f'(2) = 3*4 - 2 = 10
print(gradient) # 10.0

Visualize: Automatic Chain Rule¶

import jax
import jax.numpy as jnp

def f(x):
 a = jnp.sin(x) # a = sin(x), da/dx = cos(x)
 b = a**2 # b = a², db/da = 2a
 c = jnp.exp(b) # c = e^b, dc/db = e^b
 return c # dc/dx = dc/db * db/da * da/dx

# JAX automatically chains:
# df/dx = cos(x) * 2*sin(x) * exp(sin²(x))

grad_f = jax.grad(f)

# At x = π/4:
x = jnp.pi / 4
print(grad_f(x)) # Automatic chain rule applied!

-

Vector Functions: Jacobian and Gradient¶

Vector Input, Scalar Output¶

import jax
import jax.numpy as jnp

def loss(w):
 """w is vector [w1, w2, w3]"""
 x = jnp.array([1., 2., 3.])
 return jnp.sum((w * x)**2) # Scalar loss

w = jnp.array([1., 2., 3.])

# Gradient of scalar → gradient vector
grad_loss = jax.grad(loss)
gradient = grad_loss(w)
print(gradient.shape) # (3,) - same as w
print(gradient) # [2., 8., 18.]

Jacobian (Vector Input, Vector Output)¶

import jax
import jax.numpy as jnp

def f(x):
 """x is vector, returns vector"""
 return jnp.array([x[0]**2, x[1]**3, x[0]*x[1]])

x = jnp.array([2., 3.])

# Jacobian matrix
jacobian = jax.jacobian(f)(x)
print(jacobian.shape) # (3, 2) - 3 outputs, 2 inputs
print(jacobian) # [[4., 0.],
 # [0., 27.],
 # [3., 2.]]

Gradients of Specific Arguments¶

Single Argument¶

import jax
import jax.numpy as jnp

def f(x):
 return x**2 + 2*x + 1

# Gradient w.r.t. x
grad_f = jax.grad(f)
print(grad_f(3.0)) # 8.0

Multiple Arguments: argnums¶

import jax
import jax.numpy as jnp

def f(x, y):
 return x*y + x**2 + y**2

# Gradient w.r.t. first argument (x)
grad_f_x = jax.grad(f, argnums=0)
print(grad_f_x(2.0, 3.0)) # y + 2x = 3 + 4 = 7

# Gradient w.r.t. second argument (y)
grad_f_y = jax.grad(f, argnums=1)
print(grad_f_y(2.0, 3.0)) # x + 2y = 2 + 6 = 8

# Gradient w.r.t. both arguments
grad_f_both = jax.grad(f, argnums=(0, 1))
dx, dy = grad_f_both(2.0, 3.0)
print(dx, dy) # (7.0, 8.0)

Gradient of Arguments (None by Default)¶

import jax
import jax.numpy as jnp

def f(x, y, z):
 return x*y + jnp.sin(z)

# Only gradient w.r.t. x and z (not y)
grad_f = jax.grad(f, argnums=(0, 2))

# Returns gradient w.r.t. specified args
grad_x, grad_z = grad_f(1.0, 2.0, 0.0)
print(grad_x, grad_z) # (2.0, 1.0)

Higher-Order Derivatives¶

Second Derivative (Hessian)¶

import jax
import jax.numpy as jnp

def f(x):
 return x**4 + 3*x**2 + 2*x

# First derivative
df = jax.grad(f)

# Second derivative
d2f = jax.grad(df)

x = 2.0
print(df(x)) # 64 + 12 + 2 = 78
print(d2f(x)) # 12x² + 6 = 48 + 6 = 54

Nth Derivative¶

import jax
import jax.numpy as jnp

def f(x):
 return x**5

# Compute nth derivative
def nth_derivative(f, n, x):
 for _ in range(n):
 f = jax.grad(f)
 return f(x)

# 5th derivative of x^5 = 120
print(nth_derivative(f, 5, 2.0)) # 120.0

Hessian for Multiple Variables¶

import jax
import jax.numpy as jnp

def f(x):
 """x is vector"""
 return x[0]**2 + x[1]**2 + x[0]*x[1]

# Hessian = matrix of 2nd derivatives
hessian_f = jax.hessian(f)
x = jnp.array([1.0, 2.0])
H = hessian_f(x)
print(H) # [[2., 1.],
 # [1., 2.]]

Jacobian and Hessian¶

Jacobian-Vector Product (JVP)¶

import jax
import jax.numpy as jnp

def f(x):
 return jnp.array([x[0]**2, x[1]**3])

x = jnp.array([2., 3.])
v = jnp.array([1., 0.]) # Direction vector

# Jacobian-vector product
y, jvp = jax.linearize(f, x)
jvp_result = jvp(v)
print(jvp_result) # [4., 0.] - directional derivative

Vector-Jacobian Product (VJP)¶

import jax
import jax.numpy as jnp

def f(x):
 return jnp.array([x[0]**2 + x[1], x[0] * x[1]**2])

x = jnp.array([2., 3.])
v = jnp.array([1., 1.]) # Cotangent vector

# Vector-Jacobian product (used internally by grad)
y, vjp_fn = jax.vjp(f, x)
vjp_result = vjp_fn(v)
print(vjp_result) # Gradient direction

Practical Example: Linear Regression¶

import jax
import jax.numpy as jnp
from jax import grad

# Data
X = jnp.array([[1., 2.], [3., 4.], [5., 6.]])
y = jnp.array([1., 2., 3.])

def model(params, x):
 """Linear model: y = w·x + b"""
 w, b = params
 return jnp.dot(x, w) + b

def loss(params, X, y):
 """Mean squared error"""
 preds = jnp.array([model(params, x) for x in X])
 return jnp.mean((preds - y)**2)

# Initial parameters
params = (jnp.array([1., 1.]), 0.)

# Create gradient function
grad_loss = grad(loss)

# Training loop
learning_rate = 0.01
for step in range(100):
 grads = grad_loss(params, X, y)
 # Update: params -= lr * grads (but params are immutable!)
 w_new = params[0] - learning_rate * grads[0]
 b_new = params[1] - learning_rate * grads[1]
 params = (w_new, b_new)

 if step % 20 == 0:
 print(f"Step {step}, Loss: {loss(params, X, y):.4f}")

print(f"Learned w: {params[0]}, b: {params[1]}")

Practical Example: Neural Network¶

import jax
import jax.numpy as jnp
from jax import grad, random

def relu(x):
 return jnp.maximum(0, x)

def forward(params, x):
 """Simple 2-layer network"""
 w1, b1, w2, b2 = params
 h = relu(jnp.dot(x, w1) + b1)
 return jnp.dot(h, w2) + b2

def loss(params, x, y):
 pred = forward(params, x)
 return jnp.mean((pred - y)**2)

# Initialize
key = random.PRNGKey(0)
params = (
 random.normal(key, (2, 3)) * 0.1, # w1
 jnp.zeros(3), # b1
 random.normal(key, (3, 1)) * 0.1, # w2
 jnp.zeros(1), # b2
)

# Training
grad_loss = grad(loss)
X = random.normal(key, (10, 2))
y = random.normal(key, (10, 1))

for step in range(100):
 grads = grad_loss(params, X, y)

 # Update params
 new_params = tuple(
 p - 0.01 * g for p, g in zip(params, grads)
)
 params = new_params

print(f"Final loss: {loss(params, X, y):.4f}")

Gotchas and Debugging¶

Gotcha 1: Gradients w.r.t. Constants¶

import jax
import jax.numpy as jnp

def f(x):
 const = 5.0 # Constant (not differentiated)
 return x**2 + const

grad_f = jax.grad(f)
print(grad_f(3.0)) # 6.0 (const doesn't affect gradient)

Gotcha 2: Python Control Flow¶

import jax
import jax.numpy as jnp

def f(x):
 if x > 0: # Python control flow
 return x**2
 else:
 return -x**2

grad_f = jax.grad(f)
# Might not work or produce unexpected results!
# Use jax.lax.cond instead

# Correct
def f_correct(x):
 return jax.lax.cond(
 x > 0,
 lambda: x**2,
 lambda: -x**2
)

Gotcha 3: Gradient of Scalar Expected¶

import jax
import jax.numpy as jnp

def f(x):
 return jnp.array([x**2, x**3]) # Returns vector!

# This fails
# grad_f = jax.grad(f) # Error

# Use jacobian for vector output
jac_f = jax.jacobian(f)
print(jac_f(2.0)) # Matrix of derivatives

Gotcha 4: Tracing Issues¶

import jax
import jax.numpy as jnp

counter = 0

def f(x):
 global counter
 counter += 1 # Impure!
 return x**2

# During tracing, counter is incremented
# But gradient computation sees different behavior

Performance Tips¶

Tip 1: Jit Compile Gradient Computation¶

import jax
import jax.numpy as jnp

def loss(params, x, y):
 return jnp.mean((params @ x - y)**2)

# Slow
grad_loss = jax.grad(loss)

# Fast
grad_loss_jit = jax.jit(jax.grad(loss))

Tip 2: Reuse Gradient Functions¶

# Compile once
grad_loss_jit = jax.jit(jax.grad(loss))

# Reuse many times
for step in range(1000):
 grads = grad_loss_jit(params, X[step], y[step]) # Fast!

Summary¶

AUTOMATIC DIFFERENTIATION QUICK REFERENCE:

Basic gradient:
 grad_f = jax.grad(f)
 g = grad_f(x)

Multiple args:
 grad_f = jax.grad(f, argnums=(0, 2))
 g0, g2 = grad_f(x0, x1, x2)

Jacobian:
 J = jax.jacobian(f)(x)

Hessian:
 H = jax.hessian(f)(x)

Higher derivatives:
 d2f = jax.grad(jax.grad(f))
 d2f(x)

-

Next Steps¶

02 Just In Time Compilation Jax.Jit - Speed up gradient computation
03 Vectorization Jax.Vmap - Compute gradients for batches

Checkpoint: Implement gradient descent for y = (x-3)² (minimum at x=3)

-

Last Updated: 2026-08-09