One of JAX's unique strengths is easy higher-order differentiation. You can take the gradient of a function that itself uses gradients. This is the foundation of meta-learning algorithms like MAML (Model-Agnostic Meta-Learning).
import jax
import jax.numpy as jnp
# Second-order derivatives are trivial
def f(x):
return jnp.sin(x) * x ** 2
# First derivative
df = jax.grad(f)
print(df(1.0)) # cos(1)*1 + sin(1)*2 ≈ 2.22
# Second derivative (Hessian for scalar functions)
d2f = jax.grad(jax.grad(f))
print(d2f(1.0)) # ≈ -0.18
# Hessian for vector functions
def g(x):
return jnp.sum(x ** 3)
hessian = jax.hessian(g)
print(hessian(jnp.array([1.0, 2.0, 3.0])))
# [[6., 0., 0.],
# [0., 12., 0.],
# [0., 0., 18.]]MAML: Learning to Learn
def maml_loss(meta_params, tasks, inner_lr=0.01, inner_steps=5):
"""MAML outer loss: how well do inner-loop-adapted params perform?"""
total_loss = 0.0
for task_train, task_test in tasks:
# Inner loop: adapt to task using gradient descent
params = meta_params
for _ in range(inner_steps):
train_loss = compute_loss(params, *task_train)
grads = jax.grad(compute_loss)(params, *task_train)
params = jax.tree.map(
lambda p, g: p - inner_lr * g, params, grads)
# Outer loss: evaluate adapted params on test data
test_loss = compute_loss(params, *task_test)
total_loss += test_loss
return total_loss / len(tasks)
# The magic: differentiate through the inner loop!
meta_grads = jax.grad(maml_loss)(meta_params, tasks)
# This computes second-order gradients (gradient through gradient descent)
# In PyTorch, this requires create_graph=True and careful management.
# In JAX, it just works — grad(grad(...)) composes naturally.💡 Why This Matters
MAML requires differentiating through multiple steps of gradient descent — a second-order gradient computation. In PyTorch, this needs create_graph=True, careful memory management, and is easy to get wrong. In JAX, it's just jax.grad(outer_loss) where outer_loss internally calls jax.grad. JAX's composable transformations make higher-order differentiation almost trivially easy, which is why JAX dominates meta-learning research.
# Practical example: memory-efficient Transformer with remat + scan
def create_efficient_transformer(d_model, num_heads, d_ff, num_layers, rngs):
"""Create a transformer that uses scan + remat for efficiency."""
# Initialize one block's parameters
def init_block(rngs):
return TransformerBlock(d_model, num_heads, d_ff, rngs)
blocks = [init_block(rngs) for _ in range(num_layers)]
def forward(x):
for block in blocks:
# Remat: recompute activations in backward pass (saves memory)
x = jax.checkpoint(lambda b, x: b(x), block, x)
return x
return forward, blocks