Custom Gradient Rules: custom_vjp and custom_jvp

Sometimes the automatic gradient that JAX computes isn't what you want. Maybe you need a numerically stable gradient, a straight-through estimator, or a gradient that doesn't exist mathematically but is useful for training. JAX lets you define custom gradient rules using jax.custom_vjp (for reverse-mode) and jax.custom_jvp (for forward-mode).

custom_vjp: Custom Reverse-Mode Gradients

import jax
import jax.numpy as jnp

@jax.custom_vjp
def safe_log(x):
    """Log that's numerically stable in the forward pass."""
    return jnp.log(jnp.maximum(x, 1e-7))

def safe_log_fwd(x):
    """Forward pass: returns (output, residuals for backward)."""
    result = safe_log(x)
    return result, x  # save x for the backward pass

def safe_log_bwd(x, g):
    """Backward pass: g is the incoming gradient, x is the saved residual."""
    # Clip gradient to prevent explosion near zero
    return (g / jnp.maximum(x, 1e-7),)

safe_log.defvjp(safe_log_fwd, safe_log_bwd)

# Now gradients are clipped automatically
x = jnp.array([1.0, 0.01, 0.0001, 0.0])
grads = jax.grad(lambda x: jnp.sum(safe_log(x)))(x)
print(grads)  # [1.0, 100.0, 10000.0, 10000000.0] — capped, not inf

Straight-Through Estimator

A classic use case: quantization that uses a discrete forward pass but pretends the gradient flows straight through:

@jax.custom_vjp
def straight_through_round(x):
    """Round in forward pass, identity gradient in backward."""
    return jnp.round(x)

def ste_fwd(x):
    return straight_through_round(x), x

def ste_bwd(x, g):
    return (g,)  # gradient passes through unchanged

straight_through_round.defvjp(ste_fwd, ste_bwd)

# Forward: rounds values. Backward: gradient flows as if no rounding happened
x = jnp.array([1.3, 2.7, -0.5])
print(straight_through_round(x))  # [1., 3., -0.]
print(jax.grad(lambda x: jnp.sum(straight_through_round(x)))(x))  # [1., 1., 1.]

custom_jvp: Custom Forward-Mode Gradients

@jax.custom_jvp
def my_relu(x):
    return jnp.maximum(x, 0.0)

@my_relu.defjvp
def my_relu_jvp(primals, tangents):
    (x,) = primals
    (x_dot,) = tangents
    primal_out = my_relu(x)
    # Custom: use sigmoid as a smooth approximation of the step function
    tangent_out = x_dot * jax.nn.sigmoid(10.0 * x)
    return primal_out, tangent_out