TopolearnStable gradients for functions like log1p(exp(x)), implicit layers and straight-through estimators require custom rules.
Sometimes the derivative JAX derives by differentiating your implementation is numerically poor, inefficient, or undefined. You can supply your own rule.
`jax.custom_jvp` defines a forward-mode rule; JAX derives reverse mode by transposing it. Decorate f, then register f.defjvp(rule) where rule(primals, tangents) returns (primal_out, tangent_out). Prefer this when the rule is expressible with JAX ops, because it supports both modes and higher-order derivatives.
`jax.custom_vjp` defines reverse mode directly with two functions: fwd(*args) returns the output plus residuals to save, and bwd(residuals, g) returns a tuple with one cotangent per input. Use it for things forward-mode cannot express, like gradient clipping in the backward pass, implicit differentiation of a fixed-point solver, or wrapping code JAX cannot trace. Custom-VJP functions do not support forward-mode differentiation.
Typical use cases:
- Numerical stability:
log(1 + exp(x))overflows for large x in its naive gradient. - Implicit functions: differentiate the solution of
g(x, theta) = 0via the implicit function theorem instead of through solver iterations. - Straight-through estimators: identity gradient through
round.
Rules must return cotangents with the same structure and dtype as inputs; use None-free zero arrays for non-differentiable inputs or mark them with nondiff_argnums.
Example
import jax
import jax.numpy as jnp
@jax.custom_jvp
def softplus(x):
return jnp.logaddexp(x, 0.0)
@softplus.defjvp
def softplus_jvp(primals, tangents):
(x,), (t,) = primals, tangents
return softplus(x), jax.nn.sigmoid(x) * t # stable derivative
print(jax.grad(softplus)(100.0)) # 1.0
@jax.custom_vjp
def clip_grad(x):
return x # identity forward
fwd = lambda x: (x, None) # output, residuals (none needed)
bwd = lambda _, g: (jnp.clip(g, -1.0, 1.0),) # tuple: one cotangent per input
clip_grad.defvjp(fwd, bwd)
print(jax.grad(lambda x: 10.0 * clip_grad(x))(2.0)) # 1.0 instead of 10.0Remember
custom_jvpdefines forward mode and gets reverse mode by transpositioncustom_vjpuses fwd (save residuals) and bwd (return a tuple of cotangents)- Use them for stability, implicit differentiation and estimators like straight-through
Watch out for
- Returning a bare array instead of a tuple from the bwd function of a custom_vjp.