Chapter 8: Composing Transformations

Overview

JAX's power comes from composing transformations. This chapter covers nesting grad, jit, and vmap.

Key Topics

1. grad(grad(f)) - Second Derivatives

import jax
import jax.numpy as jnp

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

# First derivative
df = jax.grad(f)
print(df(2.0))  # 4*x³ + 6*x = 4*8 + 12 = 44

# Second derivative (Hessian for scalar)
d2f = jax.grad(df)
print(d2f(2.0))  # 12*x² + 6 = 48 + 6 = 54

2. jit(grad(f)) - Compiled Gradients

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

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

# Compile the gradient computation
grad_loss_jit = jit(grad(loss))

# Much faster for repeated calls
params = jnp.array([1., 2., 3.])
x = jnp.array([1., 2., 3.])
y = jnp.array([5., 6., 7.])

grads = grad_loss_jit(params, x, y)

3. vmap(grad(f)) - Batch Gradients

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

def loss_single(w, x, y):
    return (w * x - y)**2

# Gradient for each sample in batch
w = jnp.array([1., 2., 3.])
X = jnp.array([[1., 2., 3.],
               [4., 5., 6.]])
y = jnp.array([1., 2.])

grad_loss = vmap(grad(loss_single, argnums=0), in_axes=(None, 0, 0))
grads = grad_loss(w, X, y)

4. jit(vmap(grad(f))) - All Together

# Vectorize, differentiate, and compile - ultimate power!
f_optimized = jax.jit(jax.vmap(jax.grad(loss)))

Practical Examples

Example: Training Step Composition

def train_step(params, x, y, learning_rate):
    loss_value = loss(params, x, y)
    grads = grad(loss)(params, x, y)
    new_params = params - learning_rate * grads
    return new_params, loss_value

# Compile the entire training step
train_step_jit = jax.jit(train_step)

Best Practices

1. Order matters: jit(vmap(grad(f))) vs vmap(jit(grad(f)))
2. Profile first: Measure which composition is fastest
3. Batch when possible: vmap(grad()) for samples

Summary

• Compose transformations for power
• Understand execution order
• Profile to optimize
• Use jit outermost for compilation