When you write the rule

The Euclidean norm is three lines of arithmetic with no special functions in it. Ask for its value at the origin and you get 0.0, which is correct. Ask for its gradient there and every component comes back nan.

The algebra says why without any appeal to floating point. The derivative of the norm is x divided by the norm, and at the origin that is zero over zero. JAX is not wrong here; it evaluated the rule it was given at a point where the rule has no value.

What matters for a real model is that this is a region and not a single point. At 1e-20 the squares underflow to zero in float32, so the norm is zero while x is not, and the first component comes back inf while the others come back nan. Any padded row, masked token, or freshly zeroed embedding walks into that region on the first step.

import jax
import jax.numpy as jnp

def norm(x):
    return jnp.sqrt(jnp.sum(x ** 2))

for pt in ([0.0, 0.0, 0.0], [1e-20, 0.0, 0.0], [1e-3, 0.0, 0.0]):
    p = jnp.array(pt)
    print(norm(p).item(), jax.grad(norm)(p).tolist())

# 0.0 [nan, nan, nan]
# 0.0 [inf, nan, nan]
# 0.0010000000474974513 [1.0, 0.0, 0.0]

The reflex is to guard the output: compute the norm, and where the input was zero, return zero instead. It does exactly nothing for the gradient. jnp.where evaluates both of its branches and then selects, so the nan tangent from the branch you discarded is still multiplied by a zero from the selection, and nan times zero is nan.

The fix the FAQ prescribes is to guard the input as well, so the unsafe operation never sees the value it cannot handle. Substitute a harmless 1.0 inside the square root, then select on the same predicate outside it. Written that way the gradient at the origin is a clean zero.

It works and it is fragile in a specific way: the predicate is now written twice, and anyone editing either copy has to know that the inner one exists for the derivative rather than for the value. When the derivative is something you can state outright, saying so is the more durable repair.

import jax
import jax.numpy as jnp

def masked(x):
    s = jnp.sum(x ** 2)
    return jnp.where(s > 0.0, jnp.sqrt(s), 0.0)

def twice_masked(x):
    s = jnp.sum(x ** 2)
    safe = jnp.where(s > 0.0, s, 1.0)
    return jnp.where(s > 0.0, jnp.sqrt(safe), 0.0)

zero = jnp.zeros(3)
print(masked(zero), jax.grad(masked)(zero))
print(twice_masked(zero), jax.grad(twice_masked)(zero))

# 0.0 [nan nan nan]
# 0.0 [0. 0. 0.]

jax.custom_jvp replaces the forward rule for a function, and lesson one is why that is the load-bearing one: every other mode is built from it. The decorated rule takes the primals and tangents as two tuples and returns the primal output beside the tangent output, which for a norm is the projection of the tangent onto the unit vector, dot(x, dx) / |x|.

Two things are true about the guard inside that rule. It is written once, in the derivative, where it belongs. And it encodes a choice rather than a fact: the norm has no derivative at the origin, so returning zero there picks the subgradient of smallest length, which is the one that leaves an optimizer standing still instead of producing nan.

With the rule in place the gradient at the origin is zero, the gradient at (3, 4, 0) is the unit vector (0.6, 0.8, 0), and jax.jvp works on the same function, because a forward rule is what forward mode wanted in the first place.

import jax
import jax.numpy as jnp

@jax.custom_jvp
def norm(x):
    return jnp.sqrt(jnp.sum(x ** 2))

@norm.defjvp
def norm_jvp(primals, tangents):
    (x,), (dx,) = primals, tangents
    y = norm(x)
    safe = jnp.where(y > 0.0, y, 1.0)
    return y, jnp.where(y > 0.0, jnp.dot(x, dx) / safe, 0.0)

p = jnp.array([3.0, 4.0, 0.0])
print(jax.grad(norm)(jnp.zeros(3)))
print(jax.grad(norm)(p))
print(jax.jvp(norm, (p,), (jnp.array([1.0, 0.0, 0.0]),)))

# [0. 0. 0.]
# [0.6 0.8 0. ]
# (Array(5., dtype=float32), Array(0.6, dtype=float32))

Write the same repair as a custom_vjp and reverse mode is happy: the gradient at (3, 4, 0) is the same unit vector. Then ask for a jvp of it and the answer is a refusal in one sentence, can't apply forward-mode autodiff (jvp) to a custom_vjp function.

That refusal propagates further than it first appears. Anything that needs a forward pass over the gradient is closed off too, which includes second derivatives and every Hessian-vector product, since those are a jvp of a grad. The next lesson measures what you would have been buying there.

So the choice is not a matter of taste. Reach for custom_jvp when the derivative is something you can state going forward, and the function keeps working under every transformation and every order. Reach for custom_vjp when only the backward rule exists as a separate program: a fused kernel whose backward is its own kernel, which is what the Pallas arc and the flash-attention stage are about.

import jax
import jax.numpy as jnp

@jax.custom_vjp
def norm_v(x):
    return jnp.sqrt(jnp.sum(x ** 2))

def fwd(x):
    y = jnp.sqrt(jnp.sum(x ** 2))
    safe = jnp.where(y > 0.0, y, 1.0)
    return y, jnp.where(y > 0.0, x / safe, 0.0)

def bwd(res, ct):
    return (ct * res,)

norm_v.defvjp(fwd, bwd)

p = jnp.array([3.0, 4.0, 0.0])
print(jax.grad(norm_v)(p))
try:
    jax.jvp(norm_v, (p,), (jnp.ones(3),))
except TypeError as err:
    print(err)

# [0.6 0.8 0. ]
# can't apply forward-mode autodiff (jvp) to a custom_vjp function.

A hand-written rule is a claim, and jax.test_util.check_grads is how you test it. The docs describe it as checking gradients from automatic differentiation against finite differences, and order=2 runs the check on the derivative of the derivative as well. On the custom rule above, at (3, 4, 0), it passes.

Then look at the Hessian it produces, because there is a closed form to compare against. The exact answer is the projection away from the unit vector divided by the norm, which for this point is 0.128 and 0.072 on the diagonal, -0.096 off it, and 0.2 in the direction with no component. Those are the printed numbers, which means the rule survived being differentiated twice and stayed correct.

One guard rail comes free. A rule that returns a tangent of the wrong shape raises before it can hand you a wrong number, and the message names both the shape it expected and the shape it got.

import jax
import jax.numpy as jnp
from jax.test_util import check_grads

@jax.custom_jvp
def norm(x):
    return jnp.sqrt(jnp.sum(x ** 2))

@norm.defjvp
def norm_jvp(primals, tangents):
    (x,), (dx,) = primals, tangents
    y = norm(x)
    safe = jnp.where(y > 0.0, y, 1.0)
    return y, jnp.where(y > 0.0, jnp.dot(x, dx) / safe, 0.0)

check_grads(norm, (jnp.array([3.0, 4.0, 0.0]),), order=2)
print('order 2 ok')
print(jax.hessian(norm)(jnp.array([3.0, 4.0, 0.0])))

@jax.custom_jvp
def bad(x):
    return jnp.sum(x ** 2)

@bad.defjvp
def bad_jvp(primals, tangents):
    (x,), (dx,) = primals, tangents
    return bad(x), 2.0 * x * dx

try:
    jax.jvp(bad, (jnp.ones(3),), (jnp.ones(3),))
except TypeError as err:
    print(err)

# order 2 ok
# [[ 0.128 -0.096  0.   ]
#  [-0.096  0.072  0.   ]
#  [ 0.     0.     0.2  ]]
# Custom JVP rule must produce primal and tangent outputs with corresponding shapes
# and dtypes. Expected float32[] (tangent type of float32[]) but got float32[3].