Composing grad with jit and vmap

The real power of JAX's functional design shows when you compose transformations. Each transformation can wrap any other, creating powerful combinations with minimal code.

jit(grad(f)): Fast Gradients

import jax
import jax.numpy as jnp

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

# Compose: differentiate, then compile
fast_grad = jax.jit(jax.grad(loss_fn))

params = jnp.zeros(3)
x = jax.random.normal(jax.random.PRNGKey(0), (100, 3))
y = jax.random.normal(jax.random.PRNGKey(1), (100,))

# Fast compiled gradient computation
grads = fast_grad(params, x, y)

vmap(grad(f)): Per-Example Gradients

This is one of JAX's killer features. By combining vmap and grad, you can compute separate gradients for each example in a batch — something that's very difficult in other frameworks.

import jax
import jax.numpy as jnp

def single_loss(params, x, y):
    """Loss for ONE example."""
    pred = jnp.dot(x, params)
    return (pred - y) ** 2

# grad: gradient for one example
grad_single = jax.grad(single_loss)

# vmap(grad): gradient for each example independently
per_example_grad = jax.vmap(grad_single, in_axes=(None, 0, 0))

# jit it for speed
fast_per_example_grad = jax.jit(per_example_grad)

params = jnp.array([1.0, 2.0, 3.0])
X = jax.random.normal(jax.random.PRNGKey(0), (64, 3))  # 64 examples
y = jax.random.normal(jax.random.PRNGKey(1), (64,))

grads = fast_per_example_grad(params, X, y)
print(grads.shape)  # (64, 3) — one gradient per example!

stop_gradient: Controlling Gradient Flow

Sometimes you want part of a computation to be excluded from differentiation. jax.lax.stop_gradient treats its input as a constant during backpropagation:

import jax
import jax.numpy as jnp

def loss_with_target_network(params, target_params, x):
    """Common in RL: target network should not receive gradients."""
    pred = jnp.dot(x, params)
    # Stop gradient on target — treat as constant
    target = jax.lax.stop_gradient(jnp.dot(x, target_params))
    return jnp.mean((pred - target) ** 2)

# Gradient only flows to params, not target_params
grad_fn = jax.grad(loss_with_target_network)
grads = grad_fn(
    jnp.array([1.0, 2.0]),      # params — gets gradients
    jnp.array([0.5, 1.5]),      # target_params — no gradients
    jnp.array([[1.0, 0.5]])     # x
)
print(grads)  # Gradient with respect to params only