Chapter 4: Automatic Differentiation - jax.grad

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: Error-prone, time-consuming
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: 1M function calls!

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: ∂f_i/∂x_j
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: use JAX control flow
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: can't differentiate vector output

# ✅ 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