Chapter 10: Custom Derivatives¶

Overview¶

Sometimes you need custom gradients. JAX provides vjp, jvp for low-level control.

Topics¶

1. Custom Backward Pass with defvjp¶

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x)

# Define custom VJP
def f_vjp(x):
    y = jnp.sin(x)
    def vjp_fn(g):
        return (g * jnp.cos(x),)
    return y, vjp_fn

# Register
jax.defvjp(f, f_vjp)

2. Using custom_vjp Decorator¶

from jax import custom_vjp

@custom_vjp
def f(x):
    return jnp.sin(x)

def f_fwd(x):
    y = jnp.sin(x)
    res = (x,)  # Residuals for backward
    return y, res

def f_bwd(res, g):
    x, = res
    return (g * jnp.cos(x),)

f.defvjp(f_fwd, f_bwd)

# Now use it with grad
grad_f = jax.grad(f)
print(grad_f(1.0))  # cos(1.0) ≈ 0.540

3. Numerical Gradient Checking¶

def numerical_grad(f, x, eps=1e-5):
    grad_num = []
    for i in range(len(x)):
        x_plus = x.copy()
        x_plus[i] += eps
        x_minus = x.copy()
        x_minus[i] -= eps

        grad_i = (f(x_plus) - f(x_minus)) / (2 * eps)
        grad_num.append(grad_i)

    return jnp.array(grad_num)

# Compare with JAX
def f(x):
    return jnp.sum(x**3)

x = jnp.array([1.0, 2.0, 3.0])
grad_jax = jax.grad(f)(x)
grad_num = numerical_grad(f, x)

print("JAX gradient:", grad_jax)
print("Numerical gradient:", grad_num)
print("Close?", jnp.allclose(grad_jax, grad_num))

Summary¶

Use custom_vjp for full control
Define forward and backward passes
Always verify with numerical gradients
Useful for performance optimization