Skip to content

Chapter 10

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