Chapter 4: Advanced Automatic Differentiation
Original course: "Advanced JAX: Performance, Optimization, and Scale", Chapter 4 Advanced automatic differentiation Environment: JAX / jaxlib 0.12.0.1, CPU only, Python 3.14. Every conclusion in this notebook was verified by running it in this environment.
Learning goals
- Understand forward mode (JVP) and reverse mode (VJP) as "linearize + transpose", and use
jax.linearizeandjax.linear_transposeto take them apart. - Know the different ways to compose higher-order derivatives; get to know Taylor mode (
jax.experimental.jet), what perturbation confusion is, and why JAX doesn't suffer from it. - Choose between
jacfwdandjacrevbased on problem size, and compute sparse Jacobians efficiently withvmap(jvp)plus coloring. - Use
jax.custom_vjpfluently (including residuals,nondiff_argnums, gradient clipping, gradient reversal, STE, the implicit function theorem) andjax.custom_jvp(includingdefjvpsand automatic transposition), and verify them withcheck_grads. - Know which differentiation modes
cond,scan,fori_loopandwhile_loopeach support; know thatjax.checkpointplusscancan save backward-pass memory (covered in depth in Advanced Chapter 6 §6.6). - Understand how gradients behave for non-differentiable ops such as argmax, round, sort and top_k, how to use
stop_gradient, and common smooth approximations. - Put it together: implement an implicit layer (DEQ) and the ODE adjoint method with
custom_vjp, and cross-check against differentiating the unrolled version directly.
Sections
| Section | Original course page |
|---|---|
| 4.1 Review of forward- and reverse-mode autodiff | review-autodiff-modes |
4.2 Jacobian-vector products (JVPs) and jax.jvp | jacobian-vector-products-jvp |
4.3 Vector-Jacobian products (VJPs) and jax.vjp | vector-jacobian-products-vjp |
| 4.4 Higher-order derivatives | higher-order-derivatives |
| 4.5 Computing full Jacobians and Hessians | computing-full-jacobians-hessians |
4.6 Custom differentiation rules with jax.custom_vjp | custom-vjp-rules |
4.7 Custom differentiation rules with jax.custom_jvp | custom-jvp-rules |
| 4.8 Differentiating through control-flow primitives | differentiation-through-control-flow |
| 4.9 Handling non-differentiable functions | handling-non-differentiable-functions |
| 4.10 Practice: implementing custom gradients | practice-implementing-custom-gradient |
| Exercises / Quiz / Version notes / Summary / Cheat sheet | — |
How to read this chapter: the example code follows the structure of each course page, keeping variable names, values and code structure as close to the course as possible, while the explanations are reorganized in our own words. All code outputs come from running on this machine. Outdated or broken patterns go into ❌ cells, followed right away by the ✅ recommended version.
How this relates to the Basics part (no repeats)
- Basics 03: Automatic differentiation with grad already covered basic
jvp/vjpusage, choosing and timingjacfwd/jacrev, NaN gradients and the "double where" trick, an intro tocustom_jvp,stop_gradient, and finite-difference checks. This chapter goes deeper from there and just links back where relevant. - Basics 07: Practical topics already covered the different ways to write an HVP and their performance; Section 4.4 only includes the course's example code and doesn't repeat that discussion.
- For
vmapsee Basics 04: Automatic vectorization with vmap; forjitsee Basics 02: JIT compilation. jax.checkpoint(rematerialization) is covered in depth in Advanced Chapter 6:./06_large_scale_training.ipynb.
⚠️ Environment note: the CPU backend of jaxlib 0.12.0.1 has a known bug: once a jit closure captures a fairly large constant matrix, doing a matmul (
W @ X) followed by a reduce silently returns wrong results (for a reproduction and workaround see "Known issues in this environment" in Basics 00: Overview; for the root-cause analysis see Advanced 02 §2.3). So in this notebook, weights and data are always passed in as function arguments, and large arrays are never captured in closures.
# Common imports
import time
import inspect
from functools import partial
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.test_util import check_grads
from jax.ad_checkpoint import print_saved_residuals # shows which residuals the backward pass saves (used in Section 4.10)
print("JAX version:", jax.__version__)
print("Devices:", jax.devices())
4.1 Review of forward- and reverse-mode autodiff¶
Why automatic differentiation: automatic differentiation (AD) breaks a function into a chain of elementary operations (add, multiply, sin, exp, etc.) and applies the chain rule over and over. Finite differences suffer from truncation and rounding error, and symbolic differentiation can make expressions blow up; AD gives derivatives to machine precision at a controlled cost. jax.grad is built on the rules for these elementary operations, and underneath it maps to the two modes, forward and reverse.
Forward mode (JVP): while computing the function value, push the "sensitivity" along an input direction (the tangent vector) forward. For $f:\mathbb{R}^n\to\mathbb{R}^m$:
$$\mathrm{JVP}(f, x, v) = J_f(x)\,v = \left.\frac{d}{d\alpha} f(x+\alpha v)\right|_{\alpha=0}$$
One JVP costs about the same as one evaluation of $f(x)$, independent of the output dimension $m$. Getting the full Jacobian takes $n$ JVPs, so forward mode pays off when $n \ll m$.
Reverse mode (VJP): first do one forward evaluation and store the intermediate values you need, then propagate a cotangent vector back from the output:
$$\mathrm{VJP}(f, x, u) = u^\top J_f(x)$$
One VJP also costs about as much as evaluating $f(x)$, independent of the input dimension $n$. Getting the full Jacobian takes $m$ VJPs. For a scalar loss $L:\mathbb{R}^n\to\mathbb{R}$, $\mathrm{VJP}(L,x,1)=\nabla L(x)^\top$, which is exactly what jax.grad is built on.
Which to pick: forward mode when $n<m$; reverse mode when $m<n$. Neural network training is the classic case of the latter ($m=1$), which is why jax.grad uses reverse mode by default. The next few sections use jax.jvp and jax.vjp to control the two modes explicitly.
(Choosing between jacfwd/jacrev, and timing them, was covered in Section 3.3 of Basics 03: Automatic differentiation with grad, so we won't repeat it here.)
🔍 Extra: reverse mode = linearize + transpose¶
🟢 Must-read
Internally, JAX does not implement a separate JVP rule and VJP rule for every primitive. Here is what it does instead:
- Linearize: do a JVP of $f$ and partially evaluate it, computing up front everything that depends only on the primal values (these are the residuals). What's left is a linear function of the tangent $v \mapsto J v$.
- Transpose: the transpose of the linear function $v\mapsto Jv$ is $u \mapsto J^\top u$, which is the VJP.
So $\texttt{vjp} = \texttt{transpose} \circ \texttt{linearize}$. Each primitive only needs a JVP rule; linear primitives (one side of mul, add, dot, transpose, reduce_sum, etc.) additionally provide a transpose rule. This also explains a point we'll see in Section 4.7: as long as the tangent part of a custom_jvp is linear, it gets reverse mode automatically.
The two matching tools in the public API:
y, f_lin = jax.linearize(f, x):f_lin(v) == jvp(f,(x,),(v,))[1]. The residuals are computed once, and callingf_linmany times afterwards is cheap. For example, Krylov iterations that repeatedly compute $Jv$ at the same point are a good fit.f_T = jax.linear_transpose(f_lin, x): the function passed in must be linear; it returns its transpose.
def f(x):
return jnp.sin(x) * x**2
x = jnp.arange(3.0)
v = jnp.ones(3)
u = jnp.array([1.0, 2.0, 3.0])
# Step 1: linearize to get y = f(x) and the linear function f_lin: v -> J v
y, f_lin = jax.linearize(f, x)
print("jaxpr of f_lin (a~d are precomputed residuals such as sin x, cos x, x²; e is the tangent, and only multiply-adds on e remain):")
print(jax.make_jaxpr(f_lin)(v))
# f_lin(v) matches jvp
jvp_out = jax.jvp(f, (x,), (v,))[1]
assert jnp.allclose(f_lin(v), jvp_out)
# Step 2: transpose the linear function to get u -> J^T u
f_T = jax.linear_transpose(f_lin, x)
(ct_from_transpose,) = f_T(u)
# Exactly the same as jax.vjp: reverse mode = linearize + transpose
_, f_vjp = jax.vjp(f, x)
(ct_from_vjp,) = f_vjp(u)
print("linear_transpose(linearize):", ct_from_transpose)
print("vjp :", ct_from_vjp)
assert jnp.allclose(ct_from_transpose, ct_from_vjp)
# Cross-check against the explicit Jacobian: J^T u
J = jax.jacfwd(f)(x)
assert jnp.allclose(J.T @ u, ct_from_vjp)
# linear_transpose only accepts linear functions: transposing a nonlinear function like sin raises an error
try:
jax.linear_transpose(jnp.sin, x)(x)
raised = False
except NotImplementedError as e:
raised = True
print("NotImplementedError:", str(e)[:90])
assert raised
# The benefit of linearize: residuals are computed once, so you can vmap over many directions at the same point to batch J v
V = jnp.eye(3)
cols = jax.vmap(f_lin)(V) # each row is J e_i, i.e. column i of J
assert jnp.allclose(cols.T, J)
print("Jacobian from vmap(f_lin):\n", cols.T)
4.2 Jacobian-vector products (JVPs) and jax.jvp¶
Jacobian: the Jacobian of $f:\mathbb{R}^n\to\mathbb{R}^m$ at $x$, $J(x)=\dfrac{\partial f}{\partial x}$, is an $m\times n$ matrix. It is the best linear approximation of $f$ near $x$.
Definition of the JVP:
$$\mathrm{JVP}(x,v) = J(x)\,v = \left.\frac{d}{d\alpha} f(x+\alpha v)\right|_{\alpha=0}\in\mathbb{R}^m$$
It is the rate of change of the output when the input is perturbed infinitesimally along direction $v$. The key point is that the whole computation never builds $J(x)$ explicitly.
Signature: jax.jvp(fun, primals, tangents)
primals: a tuple (or list) of inputs;tangents: a tuple of tangents with exactly the same tree structure, shapes and dtypes asprimals. If you only want to perturb some arguments, pass zeros as the tangents of the rest (passingNonedoesn't work; see the note below);- Returns
(primal_out, tangent_out), both with the same structure.
Example 1: $f(x) = [x_0^2,\ x_0 x_1]$, evaluated at $x=[2,3]$, $v=[1,0.5]$. By hand, $J(x)=\begin{bmatrix}2x_0&0\\x_1&x_0\end{bmatrix}=\begin{bmatrix}4&0\\3&2\end{bmatrix}$, so $Jv=[4,4]$.
# Course code: a simple example
def f(x):
return jnp.array([x[0]**2, x[0] * x[1]])
x_primal = jnp.array([2.0, 3.0])
v_tangent = jnp.array([1.0, 0.5])
primal_out, tangent_out = jax.jvp(f, (x_primal,), (v_tangent,))
print(f"Primal input (x): {x_primal}")
print(f"Tangent vector (v): {v_tangent}")
print(f"Primal output f(x): {primal_out}")
print(f"Tangent output (J(x)v): {tangent_out}")
# Derived by hand: J = [[2*x0, 0], [x1, x0]] = [[4, 0], [3, 2]], J v = [4, 4]
J_manual = jnp.array([[4.0, 0.0], [3.0, 2.0]])
assert jnp.allclose(primal_out, jnp.array([4.0, 6.0]))
assert jnp.allclose(tangent_out, J_manual @ v_tangent)
Example 2 (PyTree): a linear model predict(params, inputs) = inputs @ w + b, where params is a dict {'w', 'b'}. Perturb only w (tangent of all ones), and neither b nor inputs. The tangent output should be inputs @ tangent_w + tangent_b = [30, 30].
⚠️ Note: the forward output is
inputs @ w + b = [70+0.1, 100-0.1] = [70.1, 99.9]. The second entry is 100 minus 0.1, not 100.9.⚠️ Note: don't write the tangent for
inputsasjax.lax.stop_gradient(inputs). The value ofstop_gradient(inputs)is justinputsitself, i.e.[10, 20], not a zero tangent, so the tangent output becomes[100, 130]instead of[30, 30]. PassingNonedoesn't work either; it raisesTypeErrorright away (mismatched tree structure). To say "don't perturb", passjnp.zeros_like(inputs), or fixinputswith a closure.
# ❌ Old style (outdated/wrong): use stop_gradient(inputs) as a "zero tangent", or pass None
def predict(params, inputs):
return jnp.dot(inputs, params['w']) + params['b']
params = {'w': jnp.array([[1.0, 2.0], [3.0, 4.0]]), 'b': jnp.array([0.1, -0.1])}
inputs = jnp.array([10.0, 20.0])
tangents = {'w': jnp.ones_like(params['w']), 'b': jnp.zeros_like(params['b'])}
_, t_wrong = jax.jvp(predict, (params, inputs), (tangents, jax.lax.stop_gradient(inputs)))
print("stop_gradient(inputs) as tangent ->", t_wrong, "(adds an extra contribution along inputs, inputs @ w)")
try:
jax.jvp(predict, (params, inputs), (tangents, None))
except TypeError as e:
print("Passing None ->", type(e).__name__, ":", str(e)[:80], "...")
# ✅ Recommended: pass zero tangents for arguments you don't perturb (or fix them with a closure)
primal_out, tangent_out = jax.jvp(predict, (params, inputs), (tangents, jnp.zeros_like(inputs)))
print("Inputs:", inputs)
print("Params:", params)
print("Tangents (perturbation):", tangents)
print("Primal output (prediction):", primal_out)
print("Tangent output (change in prediction due to perturbation):", tangent_out)
# Another way: jvp over params only, using a closure
_, tangent_out2 = jax.jvp(lambda p: predict(p, inputs), (params,), (tangents,))
assert jnp.allclose(primal_out, jnp.array([70.1, 99.9])) # the second entry is 99.9, not 100.9
assert jnp.allclose(tangent_out, jnp.array([30.0, 30.0]))
assert jnp.allclose(tangent_out2, tangent_out)
assert jnp.allclose(t_wrong, jnp.array([100.0, 130.0])) # stop_gradient is not a zero tangent
Uses and advantages, four in total:
- One JVP usually costs a few times one evaluation of the original function (the usual figure is 1.5–3×), independent of the output dimension;
- It is mainly used for directional derivatives, e.g. in sensitivity analysis;
- With $v=e_j$ (a one-hot vector), $J e_j$ is column $j$ of the Jacobian, and repeating this $n$ times gives the full Jacobian. In practice you just use
jax.jacfwd, which is essentiallyvmapwrapped aroundjvp; - The JVP is the foundation of JAX's autodiff; quantities like the Hessian rely on it internally (see the extra in Section 4.1: the VJP itself is obtained by linearizing the JVP and transposing it).
jax.jvp composes freely with jit and vmap, and complements the reverse mode of jax.grad/jax.vjp.
🔍 Extra: dtype rules for tangents, and float0¶
⚪ Optional: when you call
jax.jvpdirectly on a function with integer arguments
- A tangent's dtype must match its primal exactly. For example, a float32 primal with a Python
inttangent raisesTypeErrorright away. - Integer or boolean primals have no "infinitesimal perturbation", so their tangents must have the special dtype
jax.dtypes.float0. In practice, the simpler approach is to fix integer arguments with a closure orpartialso they don't appear inprimalsat all.
# Tangent dtype doesn't match the primal -> TypeError
try:
jax.jvp(jnp.sin, (jnp.float32(1.0),), (1,))
raised = False
except TypeError as e:
raised = True
print("TypeError:", str(e)[:110], "...")
assert raised
# An integer primal needs a float0 tangent
def scale_by_int(x, n):
return x * n
n = jnp.int32(3)
n_tangent = np.zeros((), dtype=jax.dtypes.float0)
_, t = jax.jvp(scale_by_int, (2.0, n), (1.0, n_tangent))
print("Derivative along x:", t)
assert t == 3.0
4.3 Vector-Jacobian products (VJPs) and jax.vjp¶
Reverse mode is the core of gradient descent in deep learning: a single backward pass gives the gradient of a scalar loss with respect to a huge number of parameters. jax.vjp is JAX's low-level function for computing VJPs, and jax.grad is a convenient wrapper around it.
Math: for $f:\mathbb{R}^n\to\mathbb{R}^m$, reverse mode computes the product of a cotangent vector and the Jacobian:
$$u^\top J_f(x)\in\mathbb{R}^{1\times n}$$
The machine-learning case: let $L=\ell(y)$ and $y=g(x)$. Then $\nabla_x L = \nabla_y L\cdot J_g(x)$, which is exactly a VJP. The cotangent starts at the output and flows backward, with initial cotangent $\partial L/\partial L = 1$.
Signature: y, vjp_fun = jax.vjp(fun, *primals)
y = fun(*primals)is the forward output;vjp_fun(u)computes $u^\top J$.umust matchyin structure, shape and dtype. The return value is a tuple with one element per primal; it's a tuple even if there's only one input.
Example 1 (scalar output): $f(x,y)=x^2\sin y$. At $x=3,\ y=\pi/2$, $\partial f/\partial x = 2x\sin y = 6$ and $\partial f/\partial y = x^2\cos y = 0$.
# Course code: scalar output
def f(x, y):
return x**2 * jnp.sin(y)
x_primal = 3.0
y_primal = jnp.pi / 2.0
y_out, vjp_fun = jax.vjp(f, x_primal, y_primal)
print("Forward output f(x, y):", y_out)
cotangent_v = 1.0
primal_grads = vjp_fun(cotangent_v)
print("Gradient w.r.t. x:", primal_grads[0]) # expected 2*x*sin(y) = 6.0
print("Gradient w.r.t. y:", primal_grads[1]) # expected x^2*cos(y) = 0.0
grad_f = jax.grad(f, argnums=(0, 1))
grads_via_grad = grad_f(x_primal, y_primal)
print("jax.grad result:", grads_via_grad)
assert jnp.allclose(primal_grads[0], 6.0)
assert jnp.allclose(primal_grads[1], 0.0, atol=1e-5)
assert all(jnp.allclose(a, b) for a, b in zip(primal_grads, grads_via_grad))
So for a scalar function, jax.grad is essentially a wrapper for "jax.vjp + an initial cotangent of 1.0".
Example 2 (non-scalar output): $g(x)=(x_0^2,\ x_1^3)$. At $x=[2,3]$, $y=[4,27]$ and $J=\mathrm{diag}(2x_0, 3x_1^2)=\mathrm{diag}(4,27)$. With cotangent $u=[10,20]$ we get $u^\top J=[40,540]$. Note that the return value is the tuple ([40., 540.],).
# Course code: non-scalar output
def g(x):
return jnp.array([x[0]**2, x[1]**3])
x_primal = jnp.array([2.0, 3.0])
y_out, vjp_fun = jax.vjp(g, x_primal)
print("y_out:", y_out) # [4., 27.]
cotangent_v = jnp.array([10.0, 20.0])
primal_grad = vjp_fun(cotangent_v)
print("u^T J:", primal_grad) # ([40., 540.],)
assert isinstance(primal_grad, tuple) and len(primal_grad) == 1
assert jnp.allclose(primal_grad[0], jnp.array([40.0, 540.0]))
Why you sometimes use jax.vjp directly, for four reasons:
- With multiple outputs, or when you need to work with intermediate gradients, using it directly is more flexible and more efficient;
- It is the basis of
jax.custom_vjp(Section 4.6); - Some optimization algorithms and AD research need direct access to the VJP;
- It helps you understand how reverse mode and
jax.gradare put together.
🔍 Extra: vjp_fun can be reused, and cotangent shapes must match exactly¶
⚪ Optional: when you use
jax.vjpdirectly, or need several rows of a Jacobian in one batch
- After
jax.vjpruns the forward pass, the residuals live in the closure ofvjp_fun. You can call the samevjp_funrepeatedly with different cotangents, and you canvmapit directly to get several rows of the Jacobian at once (that's howjacrevis implemented). - The cotangent's shape and dtype must match the output exactly, and JAX does not broadcast for you: if the output is
(3,), passing a scalar1.0raisesValueError. This is the same reasongradrequires a scalar output. jax.vjp(f, x, has_aux=True)returns(y, vjp_fun, aux), used the same way asgrad'shas_aux.
y, vjp_sq = jax.vjp(lambda x: x**2, jnp.arange(1.0, 4.0))
# The same vjp_fun, vmapped over 3 basis vectors -> 3 rows of the Jacobian
rows = jax.vmap(vjp_sq)(jnp.eye(3))[0]
assert jnp.allclose(rows, jnp.diag(2 * jnp.arange(1.0, 4.0)))
print("vmap(vjp_fun)(I) =\n", rows)
# Mismatched cotangent shapes (no automatic broadcasting)
for bad_ct in [1.0, jnp.ones(2)]:
try:
vjp_sq(bad_ct)
raised = False
except ValueError as e:
raised = True
print("ValueError:", str(e)[:100], "...")
assert raised
4.4 Higher-order derivatives¶
JAX's grad, jvp and vjp are all composable function transformations: transform a transformed function again and you get higher-order derivatives.
Scalar example: $f(x)=x^3+2x^2-3x+1$, $f'(x)=3x^2+4x-3$, $f''(x)=6x+4$, $f^{(3)}(x)=6$. At $x=2$, $f'(2)=17$ and $f''(2)=16$.
# Course code A: grad twice in a row
def f(x):
return x**3 + 2 * x**2 - 3 * x + 1
grad_f = jax.grad(f)
grad_grad_f = jax.grad(grad_f)
x_val = 2.0
first_deriv = grad_f(x_val)
second_deriv = grad_grad_f(x_val)
print("f(x) = x^3 + 2x^2 - 3x + 1")
print(f"f'(x) = 3x^2 + 4x - 3, f'({x_val}) = {first_deriv}")
print(f"f''(x) = 6x + 4, f''({x_val}) = {second_deriv}")
assert first_deriv == 17.0 and second_deriv == 16.0
Hessian-vector products (HVP): for $f:\mathbb{R}^n\to\mathbb{R}$, grad(grad(f)) doesn't work directly (grad(f) outputs a vector, which you can't grad again), and usually you don't need the full Hessian anyway, just $Hv$. There are two ways to write it:
- Forward-over-reverse: $\mathrm{jvp}(\nabla f, x, v) = J_{\nabla f}(x)\,v = H(x)\,v$;
- Reverse-over-reverse: $\nabla_x\big(\nabla f(x)\cdot v\big) = H v$.
Prefer option 1, because it needs just one backward pass with a layer of forward mode on top.
⚠️ Note: the code for option 2 is
gradovergrad, so it is reverse-over-reverse; don't call it "reverse-over-forward".
For a fuller comparison of HVP variants (including jvp(grad), grad(dot(grad, v)), vjp(grad), etc.) with timings, see Basics 07: Practical topics. Here we only include the course's example code.
# Course code B / C: two HVPs
def hvp_forward_over_reverse(f, primals, tangents):
# jvp returns (grad_f(x), H v); take the second output
return jax.jvp(jax.grad(f), primals, tangents)[1]
def hvp_reverse_over_reverse(f, primals, tangents):
x, = primals
v, = tangents
# Capture v in a lambda and take the gradient of grad(f)(x)·v once more
return jax.grad(lambda x_lambda: jnp.dot(jax.grad(f)(x_lambda), v))(x)
# Course code D: f(x, y) = x^2 y + y^3
def func(x):
return x[0]**2 * x[1] + x[1]**3
x_primal = jnp.array([1.0, 2.0])
v_tangent = jnp.array([1.0, 0.0])
hvp_val = jax.jvp(jax.grad(func), (x_primal,), (v_tangent,))[1]
# Gradient [2xy, x^2+3y^2]; Hessian [[2y, 2x], [2x, 6y]], at (1,2) = [[4,2],[2,12]], Hv = [4, 2]
print("HVP (fwd-over-rev):", hvp_val)
print("HVP (rev-over-rev):", hvp_reverse_over_reverse(func, (x_primal,), (v_tangent,)))
assert jnp.allclose(hvp_val, jnp.array([4.0, 2.0]))
assert jnp.allclose(hvp_forward_over_reverse(func, (x_primal,), (v_tangent,)),
hvp_reverse_over_reverse(func, (x_primal,), (v_tangent,)))
Full Hessian: you can use jax.hessian(f), or jacfwd(grad(f)) (column by column, JVP-based) or jacrev(grad(f)) (row by row, VJP-based).
# Course code E: three ways to get the full Hessian
hessian_matrix = jax.hessian(func)(x_primal)
hessian_jacfwd = jax.jacfwd(jax.grad(func))(x_primal)
hessian_jacrev = jax.jacrev(jax.grad(func))(x_primal)
print("jax.hessian:\n", hessian_matrix)
H_expected = jnp.array([[4.0, 2.0], [2.0, 12.0]])
for H in (hessian_matrix, hessian_jacfwd, hessian_jacrev):
assert jnp.allclose(H, H_expected)
# Course code F: third derivative
grad_grad_grad_f = jax.grad(grad_grad_f)
third_deriv = grad_grad_grad_f(x_val)
print("f^(3)(2) =", third_deriv)
assert third_deriv == 6.0
🔍 Extra: perturbation confusion, and how JAX avoids it¶
⚪ Optional: if you want to know why nested forward-mode derivatives don't go wrong (how AD is implemented)
Nesting forward mode has a classic trap. Consider this expression:
$$\frac{d}{dx}\left[x\cdot \left.\frac{d}{dy}(x+y)\right|_{y=1}\right]_{x=2}$$
The inner derivative is $\frac{d}{dy}(x+y)=1$, so the whole expression equals $\frac{d}{dx}[x]=1$.
If you implement forward AD with the most naive dual numbers, the tangents of all nesting levels get squeezed into the same "$\varepsilon$" slot. While computing the inner derivative, the outer perturbation $\varepsilon$ on $x$ gets mistaken for the inner perturbation, so the inner derivative becomes 2 and the final result is wrongly 2 as well. This is perturbation confusion.
Here is how JAX avoids it: every call to jvp/grad creates a new trace with a unique identity (its nesting level). Every tracer remembers which trace it belongs to. When an inner trace meets an outer tracer, it treats it as a constant (that isn't its own perturbation), and the outer tangent is never read by the inner one. So perturbations from different levels never mix. Below is a 30-line naive dual-number implementation that reproduces the bug, compared with JAX's result.
class Dual:
"""Naive dual numbers: all nesting levels share the same infinitesimal ε (no level tags)"""
def __init__(self, p, t):
self.p, self.t = p, t
@staticmethod
def lift(o):
return o if isinstance(o, Dual) else Dual(o, 0.0)
def __add__(self, o):
o = Dual.lift(o)
return Dual(self.p + o.p, self.t + o.t)
__radd__ = __add__
def __mul__(self, o):
o = Dual.lift(o)
return Dual(self.p * o.p, self.p * o.t + self.t * o.p)
__rmul__ = __mul__
def naive_D(f):
def df(x):
out = f(Dual(x, 1.0))
return out.t if isinstance(out, Dual) else 0.0
return df
# d/dx [ x * d/dy (x + y)|_{y=1} ] at x=2; the correct answer is 1
naive = naive_D(lambda x: x * naive_D(lambda y: x + y)(1.0))(2.0)
def jax_D(f):
return lambda x: jax.jvp(f, (x,), (1.0,))[1]
jax_result = jax_D(lambda x: x * jax_D(lambda y: x + y)(1.0))(2.0)
print("Naive dual numbers:", naive, "(perturbation confusion, wrong)")
print("JAX jvp :", jax_result, "(correct)")
assert naive == 2.0
assert jax_result == 1.0
# Nested grad is also correct: d/dx [ x * d/dy (x*y) ] = d/dx [x^2] = 2x = 4
assert jax.grad(lambda x: x * jax.grad(lambda y: x * y)(3.0))(2.0) == 4.0
🔍 Extra: Taylor mode — jax.experimental.jet¶
⚪ Optional: when you need derivatives of order 3 or higher in one go
If you nest grad $k$ times to get the $k$-th derivative, the computation graph grows quickly with $k$: each extra level differentiates the already-larger derivative graph from the previous level. Taylor mode instead pushes a truncated Taylor series forward in one pass along an input curve $x(t)=x_0 + x_1 t + \tfrac{1}{2}x_2 t^2+\dots$. Each primitive propagates the higher-order coefficients directly via Faà di Bruno's formula, and the cost grows roughly polynomially with the order.
jax.experimental.jet is still there in JAX 0.12 and imports fine. The signature is jet(fun, primals, series), returning (primal_out, series_out). Its convention: term $k$ of series is the $k$-th derivative of $x$ along the curve (not the coefficient divided by $k!$). With series=((1, 0, 0, ...),), the output is exactly $f', f'', f^{(3)}, \dots$.
Note that it lives in experimental, so the API may change. The code below also verifies that it doesn't support cond, while_loop or sort, nor custom_jvp functions (such as jax.nn.relu and jnp.logaddexp).
from jax.experimental import jet
# Check the convention: the 1st/2nd/3rd derivatives of x^3 at x=2 are 12, 12, 6
y0, (d1, d2, d3) = jet.jet(lambda x: x**3, (2.0,), ((1.0, 0.0, 0.0),))
print("jet of x^3:", y0, d1, d2, d3)
assert (y0, d1, d2, d3) == (8.0, 12.0, 12.0, 6.0)
# Cross-check against nested grad: f(x) = exp(sin x) * cos(x^2), derivatives of order 1 to 6
def g(x):
return jnp.exp(jnp.sin(x)) * jnp.cos(x**2)
K = 6
def jet_derivs(x):
_, ys = jet.jet(g, (x,), ((1.0,) + (0.0,) * (K - 1),))
return jnp.stack(ys)
def nested_grads(x):
out, h = [], g
for _ in range(K):
h = jax.grad(h)
out.append(h(x))
return jnp.stack(out)
dj, dn = jet_derivs(0.3), nested_grads(0.3)
print("jet :", dj)
print("nested grad:", dn)
assert jnp.allclose(dj, dn, rtol=1e-4)
# Compare computation-graph size: length (in characters) of the jaxpr for the K-th derivative
def kth(h, k):
for _ in range(k):
h = jax.grad(h)
return h
len_nested = len(str(jax.make_jaxpr(kth(g, K))(0.3)))
len_jet = len(str(jax.make_jaxpr(jet_derivs)(0.3)))
print(f"{K}th-order derivative, jaxpr characters: nested grad = {len_nested}, jet (gives orders 1~{K} at once) = {len_jet}")
# Primitives jet doesn't support (tested)
unsupported = {
"lax.cond": lambda x: lax.cond(x > 0, jnp.sin, jnp.cos, x),
"jnp.sort": lambda x: jnp.sort(jnp.stack([x, 2 * x])).sum(),
"lax.while_loop": lambda x: lax.while_loop(lambda c: c[0] < 3, lambda c: (c[0] + 1, c[1] * x), (0, 1.0))[1],
"jax.nn.relu(custom_jvp)": jax.nn.relu,
}
n_fail = 0
for name, fn in unsupported.items():
try:
jet.jet(fn, (0.5,), ((1.0, 0.0),))
print(f"{name:26s} supported")
except Exception as e:
n_fail += 1
print(f"{name:26s} not supported -> {type(e).__name__}")
assert n_fail == len(unsupported)
4.5 Computing full Jacobians and Hessians¶
Full $J$ and $H$ matrices are expensive to compute and to store. In machine learning you usually only need JVPs, VJPs or HVPs, but sometimes you really do need the full matrix.
Jacobian: for $f:\mathbb{R}^n\to\mathbb{R}^m$, $J_{ij}=\dfrac{\partial f_i}{\partial x_j}$, with shape $m\times n$.
jax.jacfwd: computes $Je_j$ for each input basis vector $e_j$, giving column $j$; $n$ JVPs in total, good for $n<m$;jax.jacrev: computes $e_i^\top J$ for each output basis vector $e_i$, giving row $i$; $m$ VJPs in total, good for $m<n$.jax.gradis the special case ofjacrevwith a scalar output ($m=1$);- Both use
vmapinternally to batch these JVPs/VJPs instead of a Python loop; - How to choose:
jacrevwhen $m\ll n$,jacfwdwhen $n\ll m$; when they're about the same, measure and decide.
Example function: $f(x)=[x_0^2x_1,\ \sin x_2]$, i.e. $\mathbb{R}^3\to\mathbb{R}^2$. At $x=[1,2,\pi/2]$, $J=\begin{bmatrix}2x_0x_1 & x_0^2 & 0\\ 0&0&\cos x_2\end{bmatrix}=\begin{bmatrix}4&1&0\\0&0&\approx 0\end{bmatrix}$.
# Course code A / B: jacfwd and jacrev
def func(x):
return jnp.array([x[0]**2 * x[1], jnp.sin(x[2])])
x_in = jnp.array([1.0, 2.0, jnp.pi / 2])
jacobian_fwd = jax.jacfwd(func)(x_in)
jacobian_rev = jax.jacrev(func)(x_in)
print("Input:", x_in, " output:", func(x_in))
print("jacfwd (shape %s):\n" % (jacobian_fwd.shape,), jacobian_fwd)
J_expected = jnp.array([[4.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
assert jacobian_fwd.shape == (2, 3)
assert jnp.allclose(jacobian_fwd, J_expected, atol=1e-6)
assert jnp.allclose(jacobian_rev, jacobian_fwd)
# Course code C: build it by hand with vmap + jvp (each input basis vector -> one column)
n = x_in.shape[0]
basis_vectors_in = jnp.eye(n)
primals_out, jac_cols = jax.vmap(lambda tangent: jax.jvp(func, (x_in,), (tangent,)),
in_axes=0)(basis_vectors_in)
jacobian_vmap_jvp = jac_cols.T # vmap stacks the columns along axis 0, so transpose
# Course code D: build it by hand with vmap + vjp (each output basis vector -> one row)
primals_out, vjp_fun = jax.vjp(func, x_in)
m = primals_out.shape[0]
basis_vectors_out = jnp.eye(m)
jac_rows = jax.vmap(vjp_fun, in_axes=0)(basis_vectors_out)
jacobian_vmap_vjp = jac_rows[0] # vjp_fun returns a tuple; take the first element
assert jnp.allclose(jacobian_vmap_jvp, jacobian_fwd)
assert jnp.allclose(jacobian_vmap_vjp, jacobian_rev)
print("Jacobians built with vmap(jvp) and vmap(vjp) match jacfwd/jacrev")
Hessian: $H_{ij}=\dfrac{\partial^2 f}{\partial x_i\,\partial x_j}$, with shape $n\times n$. When the second derivatives are continuous, $H$ is symmetric. Using jacfwd on the outside is usually better, because grad(f) is a "square" function $\mathbb{R}^n\to\mathbb{R}^n$ and forward mode has lower constant overhead.
⚠️ Note: how is
jax.hessian(f)actually implemented? Some say "jax.jacfwd(jax.grad(f))", others say "jacfwd(jacrev(f))orjacrev(jacfwd(f))". Looking at the JAX 0.12 source, it is alwaysjacfwd(jacrev(f)), i.e. forward-over-reverse. (For a scalar $f$,jacrev(f)andgrad(f)compute the same thing.)
Performance: a Hessian has $O(n^2)$ entries and needs $O(n^2)$ storage. Optimization algorithms usually only use HVPs, and one HVP costs about two gradient computations. Cases that really need the full matrix: Newton's method with an explicit linear solve, very low-dimensional problems, and debugging and analysis.
# Look at the source of jax.hessian
src = inspect.getsource(jax.hessian)
print(src[src.index("return jacfwd"):].strip())
assert "jacfwd(jacrev(" in src
# Course code E: jax.hessian
def scalar_func(x):
return x[0]**2 * x[1] + x[1]**3
x_in = jnp.array([1.0, 2.0])
hessian_matrix = jax.hessian(scalar_func)(x_in)
print("Hessian:\n", hessian_matrix)
assert hessian_matrix.shape == (2, 2)
assert jnp.allclose(hessian_matrix, jnp.array([[4.0, 2.0], [2.0, 12.0]]))
Course code F compares "build the full Hessian, then multiply by $v$" with "get the HVP directly with a JVP". $H=\begin{bmatrix}4&2\\2&12\end{bmatrix}$ and $v=[0.5,-0.5]$, so $Hv=[1,-5]$.
⚠️ Note: if method 2 is written as
jax.jvp(grad_f, (x_in,), (v_)), there are two problems:v_is an undefined variable (the one defined isv), so it raisesNameError; and(v)is not a tuple. The correct form is(v,).
# ❌ Old style (outdated/wrong): the tangent is written as (v_), which is undefined and not a tuple
v = jnp.array([0.5, -0.5])
grad_f = jax.grad(scalar_func)
try:
_, hvp_efficient = jax.jvp(grad_f, (x_in,), (v_))
except NameError as e:
print("NameError:", e)
# ✅ Recommended
v = jnp.array([0.5, -0.5])
# Method 1: explicit Hessian
hessian_matrix = jax.hessian(scalar_func)(x_in)
hvp_explicit = hessian_matrix @ v
# Method 2: jvp(grad) gives H v directly, without building H
grad_f = jax.grad(scalar_func)
_, hvp_efficient = jax.jvp(grad_f, (x_in,), (v,))
print("Explicit H @ v:", hvp_explicit)
print("jvp(grad) :", hvp_efficient)
assert jnp.allclose(hvp_explicit, jnp.array([1.0, -5.0]))
assert jnp.allclose(hvp_efficient, hvp_explicit)
🔍 Extra: sparse / structured Jacobians and coloring¶
⚪ Optional: when you need a large sparse (e.g. banded) Jacobian
jacfwd does one JVP per basis vector, $n$ in total. If $J$ is sparse, say tridiagonal, that's wasteful: column $j$ only has nonzeros in rows $j-1, j, j+1$.
The coloring idea (Curtis–Powell–Reid): if two columns have no nonzero rows in common, they can "share" one JVP: just add their two basis vectors and use that as the tangent, because their nonzeros won't collide in the result. For a tridiagonal matrix, split the columns into 3 groups by $j \bmod 3$ (3 colors), with seed vector $s_c = \sum_{j\equiv c\ (\mathrm{mod}\ 3)} e_j$ for each group. Then
$$J s_c = \sum_{j\equiv c} J_{:,j}$$
For row $i$, only one $j\equiv c$ satisfies $|i-j|\le 1$, so the entries can be decompressed without ambiguity: $J_{ij} = (J s_{j \bmod 3})_i$ (when $|i-j|\le1$).
That way, no matter how large $n$ is, you need only 3 JVPs (batched with vmap), while jacfwd needs $n$. Reverse mode works the same way: color the rows, then use vmap(vjp).
def f_band(x):
"""f_i = x_i^2 - x_{i-1} x_{i+1} + sin(x_i); the Jacobian is tridiagonal"""
xp = jnp.concatenate([x[1:], jnp.zeros(1, x.dtype)])
xm = jnp.concatenate([jnp.zeros(1, x.dtype), x[:-1]])
return x**2 - xm * xp + jnp.sin(x)
def colored_jacobian_tridiag(f, x, n_colors=3):
n = x.shape[0]
j = jnp.arange(n)
seeds = (j[None, :] % n_colors == jnp.arange(n_colors)[:, None]).astype(x.dtype) # (3, n)
compressed = jax.vmap(lambda s: jax.jvp(f, (x,), (s,))[1])(seeds) # (3, n): only 3 JVPs
ii, jj = j[:, None], j[None, :]
return jnp.where(jnp.abs(ii - jj) <= 1, compressed[jj % n_colors, ii], 0.0)
x9 = jnp.linspace(0.1, 1.0, 9)
J_dense = jax.jacfwd(f_band)(x9)
J_color = colored_jacobian_tridiag(f_band, x9)
print("Tridiagonal Jacobian (coloring, 3 JVPs):\n", np.round(np.asarray(J_color), 3))
assert jnp.allclose(J_color, J_dense)
# Timing comparison (n = 3000): jacfwd needs 3000 JVPs and outputs a dense n×n matrix;
# coloring does only 3 JVPs and outputs a compressed 3×n form (it already holds every nonzero; for a tridiagonal structure there's no need to expand it to dense)
def compressed_tridiag(f, x, n_colors=3):
n = x.shape[0]
j = jnp.arange(n)
seeds = (j[None, :] % n_colors == jnp.arange(n_colors)[:, None]).astype(x.dtype)
return jax.vmap(lambda s: jax.jvp(f, (x,), (s,))[1])(seeds)
n_big = 3000
x_big = jnp.linspace(0.1, 1.0, n_big)
jac_dense_fn = jax.jit(jax.jacfwd(f_band))
jac_comp_fn = jax.jit(lambda x: compressed_tridiag(f_band, x))
# Warm up (compile)
jac_dense_fn(x_big).block_until_ready()
jac_comp_fn(x_big).block_until_ready()
def bench(fn, x, reps=10):
t0 = time.perf_counter()
for _ in range(reps):
fn(x).block_until_ready()
return (time.perf_counter() - t0) / reps * 1e3
t_dense, t_comp = bench(jac_dense_fn, x_big), bench(jac_comp_fn, x_big)
print(f"jacfwd (dense {n_big}×{n_big}): {t_dense:.2f} ms")
print(f"coloring (compressed 3×{n_big}): {t_comp:.2f} ms (ratio on this machine about {t_dense / t_comp:.1f}×)")
# Correctness: the diagonals in the compressed form match the dense Jacobian
Jd, C = jac_dense_fn(x_big), jac_comp_fn(x_big)
idx = jnp.arange(n_big)
assert jnp.allclose(C[idx % 3, idx], jnp.diag(Jd)) # main diagonal J[i,i]
assert jnp.allclose(C[(idx[1:]) % 3, idx[:-1]], jnp.diag(Jd, k=1)) # superdiagonal J[i,i+1]
assert jnp.allclose(C[(idx[:-1]) % 3, idx[1:]], jnp.diag(Jd, k=-1)) # subdiagonal J[i+1,i]
Note: if you decompress the result back into a dense $n\times n$ matrix (like colored_jacobian_tridiag above), just writing the dense matrix costs $O(n^2)$. Measured on this machine, coloring is then no faster than jacfwd. The benefit of coloring comes from computing and storing only the nonzeros.
Practical tip: the key to coloring is knowing the sparsity pattern in advance. If you don't, you can run jacfwd once at a few random points to get the nonzero structure, or analyze dependencies in the jaxpr, and then find the number of colors with greedy graph coloring. In truly large-scale settings, store the decompressed result in a sparse format (e.g. jax.experimental.sparse.BCOO) instead of writing out a dense matrix as above.
4.6 Custom differentiation rules with jax.custom_vjp¶
Common reasons to define a custom VJP: numerical stability, performance (you already know a cheaper gradient formula), mathematical simplification, implicit functions (e.g. fixed points, solutions of optimization problems), and approximate gradients (STE, gradient clipping, etc.).
Structure:
@jax.custom_vjp
def f(x, y): ... # the original function: called directly when not differentiating
def f_fwd(x, y): # forward: same signature as f
return out, residuals # returns (output, intermediates needed by the backward pass)
def f_bwd(residuals, g): # backward: g is the output cotangent, same structure as out
return (gx, gy) # a tuple, length = number of positional arguments of f
f.defvjp(f_fwd, f_bwd)
In reverse mode, jax.grad does not differentiate the Python code of f. Instead it runs f_fwd first to save the residuals, and calls f_bwd on the way back.
Example 1: $f(x,y)=x\sin y$, with $\partial f/\partial x=\sin y$ and $\partial f/\partial y=x\cos y$. The backward rule is $\bar x = g\sin y$, $\bar y = g\,x\cos y$. Test with $L = f^2$: at $x=2,\ y=\pi/2$, $f=2$, $L=4$, $\partial L/\partial x = 2f\sin y = 4$, and $\partial L/\partial y = 2f\,x\cos y\approx 0$.
⚠️ Note: this example doesn't use random numbers, so there's no need to create
key = jax.random.PRNGKey(0). By the way,PRNGKeyreturns an old-style rawuint32[2]key; when you do need random numbers, usejax.random.key(which returns a typed key). (For an intro to keys see Basics 01: Introduction to JAX; for deeper random-number management see Advanced 10: Random numbers and sharding.)
# ❌ Old style (outdated/wrong): creates a jax.random.PRNGKey(0) that is never used (old-style raw uint32[2] key)
key_old = jax.random.PRNGKey(0)
print("PRNGKey:", key_old.dtype, key_old.shape, "(old style, and this example doesn't need random numbers at all)")
# ✅ Recommended: use jax.random.key (typed key) when you need random numbers; this example doesn't, so just drop it
key_new = jax.random.key(0)
print("jax.random.key:", key_new.dtype, key_new.shape)
assert jnp.issubdtype(key_new.dtype, jax.dtypes.prng_key)
# Course code B: custom VJP for f(x, y) = x * sin(y)
@jax.custom_vjp
def f(x, y):
return x * jnp.sin(y)
def f_fwd(x, y):
primals_out = f(x, y) # calling f itself inside fwd is allowed (the custom rule isn't used again here)
residuals = (x, y, jnp.sin(y), jnp.cos(y)) # this saves 4 values; (x, sin y, cos y) would actually be enough
return primals_out, residuals
def f_bwd(residuals, cotangents_in):
x, y, sin_y, cos_y = residuals
g = cotangents_in
grad_x = g * sin_y
grad_y = g * x * cos_y
return (grad_x, grad_y)
f.defvjp(f_fwd, f_bwd)
x_val = jnp.float32(2.0)
y_val = jnp.pi / 2.0
value, grads = jax.value_and_grad(lambda x, y: f(x, y)**2, argnums=(0, 1))(x_val, y_val)
value_auto, grads_auto = jax.value_and_grad(lambda x, y: (x * jnp.sin(y))**2, argnums=(0, 1))(x_val, y_val)
print("Custom:", value, grads)
print("Auto :", value_auto, grads_auto)
assert jnp.allclose(value, 4.0) and jnp.allclose(grads[0], 4.0) and jnp.allclose(grads[1], 0.0, atol=1e-5)
assert all(jnp.allclose(a, b, atol=1e-6) for a, b in zip(grads, grads_auto))
Example 2 (a stable gradient for $\log(1+e^x)$): $f(x)=\log(1+e^x)$, $f'(x)=\dfrac{e^x}{1+e^x}=\dfrac{1}{1+e^{-x}}=\sigma(x)$. The idea: keep log1p(exp(x)) in the forward pass, but use $\sigma(x)$ in the backward pass, so the gradient doesn't break for large $x$.
⚠️ Note: written this way, only the gradient is stable; the forward pass isn't. The forward is still
log1p(exp(x)), and in float32 at $x=100$ theexpoverflows, so the forward output isinf. A truly stable forward pass isjnp.logaddexp(x, 0.0)(which is howjax.nn.softplusis implemented).⚠️ Note: don't expect
try/except FloatingPointErrorto catch "the unstable gradient overflowing at $x=100$". By default JAX raises no exception at all; the unstable gradient silently returnsnan($\infty/\infty$ shows up in the backward pass), so theexceptbranch never runs. Only withjax.debug_nans(True)turned on do you get aFloatingPointError. Also note a known issue in this environment: oncevmaphas been traced in the process,jax_debug_nansmay miss NaNs; see Advanced 08: Debugging toolchain §3.1.
# ❌ Old style (outdated/wrong): the forward pass still overflows; try/except FloatingPointError catches nothing
@jax.custom_vjp
def stable_log1pexp(x):
return jnp.log1p(jnp.exp(x))
def stable_log1pexp_fwd(x):
exp_x = jnp.exp(x)
output = jnp.log1p(exp_x)
return output, (x,)
def stable_log1pexp_bwd(residuals, cotangent_in):
(x,) = residuals
return (cotangent_in * (1.0 / (1.0 + jnp.exp(-x))),)
stable_log1pexp.defvjp(stable_log1pexp_fwd, stable_log1pexp_bwd)
def unstable_log1pexp(x):
return jnp.log(1.0 + jnp.exp(x))
grad_stable = jax.grad(stable_log1pexp)
grad_unstable = jax.grad(unstable_log1pexp)
x_small, x_problematic, x_large = jnp.float32(1.0), jnp.float32(35.0), jnp.float32(100.0)
print(f"x=1 stable grad {grad_stable(x_small):.6f} unstable grad {grad_unstable(x_small):.6f}")
print(f"x=35 stable grad {grad_stable(x_problematic):.6f} unstable grad {grad_unstable(x_problematic):.6f}")
print(f"x=100 stable grad {grad_stable(x_large)} but the forward stable_log1pexp(100) = {stable_log1pexp(x_large)}")
caught = False
try:
g_bad = grad_unstable(x_large)
print("No exception raised; unstable grad =", g_bad)
except FloatingPointError:
caught = True
print("Did the except branch run:", caught)
# ✅ Recommended: logaddexp in the forward pass (stable), sigmoid in the backward pass
@jax.custom_vjp
def log1pexp(x):
return jnp.logaddexp(x, 0.0)
def log1pexp_fwd(x):
return log1pexp(x), x
def log1pexp_bwd(x, g):
return (g * jax.nn.sigmoid(x),)
log1pexp.defvjp(log1pexp_fwd, log1pexp_bwd)
print("log1pexp(100) =", log1pexp(100.0), " grad =", jax.grad(log1pexp)(100.0))
print("gradient of built-in jax.nn.softplus =", jax.grad(jax.nn.softplus)(100.0))
assert jnp.isinf(stable_log1pexp(x_large)) # the ❌ version's forward pass overflows
assert log1pexp(100.0) == 100.0 # the fixed version's forward pass is correct
assert not caught and jnp.isnan(grad_unstable(x_large)) # no exception, silently returns NaN
assert jax.grad(log1pexp)(100.0) == 1.0
# To make NaNs raise right away, use debug_nans (for debugging; it has a performance cost)
with jax.debug_nans(True):
try:
grad_unstable(x_large)
raised = False
except FloatingPointError:
raised = True
print("Does debug_nans raise FloatingPointError:", raised)
assert raised
A few things to keep in mind:
- Correctness: always cross-check against numerical differentiation or JAX's own autodiff (see
check_gradsbelow); - Store only what you need in residuals: they stay in memory until the backward pass finishes;
- Functions with a custom VJP compose with
jit,vmapand nested reverse-mode differentiation (but not forward mode; see Section 4.7); - If
fhas arguments that don't need gradients,f_bwdreturnsNonein those positions.
⚠️ Note: about the last point: returning
Nonefor an input inbwddoes work; JAX treats it as a zero gradient. But if the argument is not an array (a Python function, a string, a config object), you must declare it withnondiff_argnums; see the extra below.
# If bwd returns None for an input, JAX treats it as a zero gradient
@jax.custom_vjp
def scaled(x, s):
return x * s
scaled.defvjp(lambda x, s: (x * s, s), lambda s, g: (g * s, None))
gx, gs = jax.grad(scaled, argnums=(0, 1))(2.0, 3.0)
print("Gradient w.r.t. x:", gx, " gradient w.r.t. s (bwd returns None):", gs)
assert gx == 3.0 and gs == 0.0
🔍 Extra: signature details, residuals, and nondiff_argnums¶
🟢 Must-read
| Case | fwd signature | bwd signature | bwd returns |
|---|---|---|---|
| Plain | fwd(*args) | bwd(res, g) | A tuple of length len(args) |
nondiff_argnums=(0,) | fwd(*args) (still all arguments) | bwd(nondiff0, res, g) (non-differentiable arguments come first) | A tuple with only the differentiable arguments |
Key points:
- The output of
fwdmust have the same structure as the output off;residualscan be any pytree (tuple, dict,None). nondiff_argnumsis meant for non-array arguments, such as Python functions, integer hyperparameters or dtypes. These values can't be tracers: if you pass a traced array in anondiff_argnumsposition insidejit, you getUnexpectedTracerError. Array-valued "don't differentiate" arguments should be passed as normal arguments, withbwdreturningNoneor zeros for them.- By default (
symbolic_zeros=False), thegthatbwdreceives is always a real array: if some output doesn't contribute to the loss, the correspondinggis a zero array, notNone(verified in Section 4.10).
Below we use nondiff_argnums to implement gradient clipping and a gradient reversal layer, and verify the trap from point 2.
# Gradient clipping: identity in the forward pass, clip the cotangent to [lo, hi] in the backward pass
@partial(jax.custom_vjp, nondiff_argnums=(0, 1))
def clip_gradient(lo, hi, x):
return x
def clip_gradient_fwd(lo, hi, x):
return x, None # no residuals needed
def clip_gradient_bwd(lo, hi, res, g): # non-differentiable arguments come first
return (jnp.clip(g, lo, hi),) # return a gradient only for the differentiable argument x
clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)
loss = lambda x: jnp.sum(10.0 * clip_gradient(-1.0, 1.0, x) ** 2)
xs = jnp.array([-2.0, 0.01, 0.03, 3.0])
print("Original gradient:", jax.grad(lambda x: jnp.sum(10.0 * x**2))(xs))
print("Clipped gradient :", jax.grad(loss)(xs))
assert jnp.allclose(jax.grad(loss)(xs), jnp.clip(20.0 * xs, -1.0, 1.0))
assert jnp.allclose(clip_gradient(-1.0, 1.0, xs), xs) # identity in the forward pass
# Gradient Reversal Layer (DANN): identity in the forward pass, multiply by -λ in the backward pass
@jax.custom_vjp
def grad_reverse(x, lam):
return x
def grad_reverse_fwd(x, lam):
return x, lam
def grad_reverse_bwd(lam, g):
return (-lam * g, None) # λ is an array argument (can be traced by jit) with no gradient, so return None
grad_reverse.defvjp(grad_reverse_fwd, grad_reverse_bwd)
# Feature extractor h = w*x followed by a "domain classifier" loss = h^2; the GRL flips the sign of w's gradient
def domain_loss(w, x, lam):
h = w * x
return grad_reverse(h, lam) ** 2
w0, x0, lam0 = 1.5, 2.0, 0.5
g_plain = jax.grad(lambda w: (w * x0) ** 2)(w0)
g_grl = jax.jit(jax.grad(domain_loss))(w0, x0, lam0)
print("Plain gradient:", g_plain, " GRL gradient:", g_grl)
assert jnp.allclose(g_grl, -lam0 * g_plain)
# Trap: putting a jit-traced array in nondiff_argnums -> UnexpectedTracerError
@partial(jax.custom_vjp, nondiff_argnums=(1,))
def grl_bad(x, lam):
return x
grl_bad.defvjp(lambda x, lam: (x, None), lambda lam, r, g: (-lam * g,))
print("Works when lam is a Python constant:", jax.grad(lambda x: grl_bad(x, 0.5))(1.0))
try:
jax.jit(jax.grad(lambda x, lam: grl_bad(x, lam)))(1.0, 0.5)
raised = False
except jax.errors.UnexpectedTracerError as e:
raised = True
print("UnexpectedTracerError:", str(e)[:100], "...")
assert raised
The straight-through estimator (STE) is another classic case of "specifying the gradient by hand". It's covered in Section 4.9, which gives both a custom_vjp version and a stop_gradient version.
🔍 Extra: the implicit function theorem — differentiating fixed-point iterations¶
⚪ Optional: when you need to differentiate fixed-point iterations, optimization layers or implicit layers (DEQ)
Many quantities are defined as the result of an iterative solve, such as a fixed point $x^\star = T(a, x^\star)$, the optimum of an optimization problem, or a root of $F(a,x)=0$. Differentiating the iteration directly has three problems:
- You have to save every step's intermediates, so memory is $O(\text{number of iterations})$;
- How accurate the gradient is depends on whether the iteration converged well enough;
- If it's implemented with
while_loop(number of iterations depends on convergence), reverse-mode differentiation is impossible (Section 4.8).
The implicit function theorem: differentiate both sides of $x^\star = T(a,x^\star)$ to get $\dfrac{\partial x^\star}{\partial a} = \Big(I-\dfrac{\partial T}{\partial x}\Big)^{-1}\dfrac{\partial T}{\partial a}$. In reverse mode, given a cotangent $g$:
$$\bar a = u^\top\frac{\partial T}{\partial a},\qquad u = g + \Big(\frac{\partial T}{\partial x}\Big)^\top u$$
Note that $u$ is itself a fixed-point problem, which can also be solved iteratively with while_loop; it converges as long as $T$ is a contraction in $x$. The whole backward pass only needs the final solution $x^\star$ as a residual, so memory is $O(1)$.
@partial(jax.custom_vjp, nondiff_argnums=(0,))
def fixed_point(T, a, x0):
"""Find the fixed point of x = T(a, x) with while_loop (the number of iterations depends on convergence)"""
def cond(c):
x, x_prev, i = c
return (jnp.max(jnp.abs(x - x_prev)) > 1e-6) & (i < 500)
def body(c):
x, _, i = c
return T(a, x), x, i + 1
return lax.while_loop(cond, body, (T(a, x0), x0, 0))[0]
def fixed_point_fwd(T, a, x0):
x_star = fixed_point(T, a, x0)
return x_star, (a, x_star) # save only the solution itself
def fixed_point_bwd(T, res, g):
a, x_star = res
_, vjp_a = jax.vjp(lambda a: T(a, x_star), a)
_, vjp_x = jax.vjp(lambda x: T(a, x), x_star)
# Adjoint fixed point u = g + (∂T/∂x)^T u, also iterated with while_loop
def cond(c):
u, u_prev, i = c
return (jnp.max(jnp.abs(u - u_prev)) > 1e-7) & (i < 500)
def body(c):
u, _, i = c
return g + vjp_x(u)[0], u, i + 1
u = lax.while_loop(cond, body, (g + vjp_x(g)[0], g, 0))[0]
(a_bar,) = vjp_a(u)
return a_bar, jnp.zeros_like(x_star) # the solution doesn't depend on the initial value, so x0's gradient is 0
fixed_point.defvjp(fixed_point_fwd, fixed_point_bwd)
# Example 1: Newton iteration for sqrt(a): x = (x + a/x)/2, analytic derivative d sqrt(a)/da = 1/(2 sqrt(a))
newton_T = lambda a, x: 0.5 * (x + a / x)
a = 2.0
x_star = fixed_point(newton_T, a, 1.0)
g_implicit = jax.grad(lambda a: fixed_point(newton_T, a, 1.0))(a)
def unrolled_newton(a, n_iter=30): # unrolled version: fixed number of steps, plain backprop
return lax.fori_loop(0, n_iter, lambda i, x: newton_T(a, x), 1.0)
g_unrolled = jax.grad(unrolled_newton)(a)
print(f"x* = {x_star:.7f} (sqrt(2) = {np.sqrt(2):.7f})")
print(f"implicit grad {g_implicit:.7f} | unrolled grad {g_unrolled:.7f} | analytic {0.5 / np.sqrt(2):.7f}")
assert jnp.allclose(x_star, np.sqrt(2.0))
assert jnp.allclose(g_implicit, 0.5 / np.sqrt(2.0))
assert jnp.allclose(g_implicit, g_unrolled)
# Example 2: x = cos(a x), analytic: dx/da = -x sin(ax) / (1 + a sin(ax))
cos_T = lambda a, x: jnp.cos(a * x)
a = 1.0
x_star = fixed_point(cos_T, a, 0.5)
g_implicit = jax.jit(jax.grad(lambda a: fixed_point(cos_T, a, 0.5)))(a) # can be jitted
g_analytic = -x_star * jnp.sin(a * x_star) / (1 + a * jnp.sin(a * x_star))
print(f"x* = {x_star:.6f}, implicit grad = {g_implicit:.6f}, analytic = {g_analytic:.6f}")
assert jnp.allclose(g_implicit, g_analytic, atol=1e-5)
# For comparison: reverse-mode differentiation of the plain while_loop version fails (details in Section 4.8)
def fixed_point_plain(a):
return lax.while_loop(lambda c: jnp.abs(c[0] - c[1]) > 1e-6,
lambda c: (cos_T(a, c[0]), c[0]), (cos_T(a, 0.5), 0.5))[0]
try:
jax.grad(fixed_point_plain)(a)
raised = False
except ValueError as e:
raised = True
print("Direct grad(while_loop):", type(e).__name__, "-", str(e)[:75], "...")
assert raised
🔍 Extra: verifying custom rules with jax.test_util.check_grads¶
🟢 Must-read
check_grads(f, args, order, modes=('fwd','rev'), atol, rtol, eps) cross-checks the JVP/VJP against finite differences. fwd mode compares the JVP directly with a numerical JVP; rev mode picks random directions $t$ and $c$ and checks $\langle t,\ \mathrm{VJP}(c)\rangle \approx \langle \text{numerical JVP}(t),\ c\rangle$, i.e. that the VJP is the transpose of the numerical derivative. Both modes only check along random directions; they never compare JAX's own JVP and VJP directly with each other. With order=2 it also checks second derivatives. In JAX 0.12 it still lives in jax.test_util. A few lessons:
custom_vjpfunctions can only usemodes=['rev'], because they have no forward mode;- The default tolerances are tight in float32; loosen
atol/rtolfor iterative solvers; - It catches a "wrongly written bwd". Rules that deliberately differ from the true derivative, like gradient clipping, GRL and STE, of course fail the check. That's expected, and it's a handy way to confirm the rule is actually taking effect.
# 1) Correct custom rules: pass
check_grads(f, (jnp.float32(2.0), jnp.float32(0.7)), order=1, modes=["rev"])
check_grads(log1pexp, (jnp.float32(3.0),), order=1, modes=["rev"])
check_grads(lambda a: fixed_point(cos_T, a, 0.5), (jnp.float32(1.0),), order=1, modes=["rev"], atol=1e-2, rtol=1e-2)
print("f / log1pexp / fixed_point pass check_grads")
# 2) A deliberately wrong bwd: caught
@jax.custom_vjp
def buggy_sin(x):
return jnp.sin(x)
buggy_sin.defvjp(lambda x: (jnp.sin(x), x), lambda x, g: (g * jnp.sin(x),)) # should be cos
try:
check_grads(buggy_sin, (jnp.float32(1.0),), order=1, modes=["rev"])
caught_bug = False
except AssertionError:
caught_bug = True
print("check_grads caught the wrong bwd:", caught_bug)
assert caught_bug
# 3) A rule that deliberately changes the gradient (GRL) also "fails", which shows the rule is working
try:
# Note: pass jnp scalars. In this example, passing a Python float 1.0 makes check_grads raise AttributeError when it calls .reshape internally (tested)
check_grads(lambda x: grad_reverse(x, 0.5), (jnp.float32(1.0),), order=1, modes=["rev"])
grl_passes = True
except AssertionError:
grl_passes = False
assert not grl_passes
4.7 Custom differentiation rules with jax.custom_jvp¶
jax.grad/jax.vjp are reverse mode. Sometimes we want to control the forward-mode rule directly, i.e. how a given tangent flows to the output. Four typical reasons for a custom JVP:
- Numerical stability:
sqrtandlognear 0; - Performance;
- Conventions at non-differentiable points: e.g. the boundaries of
jnp.clip; - External code: e.g.
jax.pure_callbackhas no differentiation rule of its own.
How to write it: decorate the original function with @jax.custom_jvp, then register the rule with @original_function.defjvp. The rule receives primals (a tuple, e.g. (x,)) and tangents (a tuple, e.g. (t,)) and returns (primal_out, tangent_out); mathematically it computes $J(x)\,t$.
The example is a "stable sqrt" whose goal is a finite derivative at $x=0$ (by convention 0; the true derivative $\frac{1}{2\sqrt{x}}\to\infty$).
⚠️ Note: the
safe_x = where(x==0, 1, x)trick belongs only in the derivative formula, not in the forward pass. If the forward pass also usessafe_x,stable_sqrt(0.0)returns 1.0 instead of 0, so the forward result is wrong (that's what happens in the ❌ cell below). Just usejnp.sqrt(x)in the forward pass; $\sqrt 0 = 0$ is perfectly fine.
# ❌ Old style (outdated/wrong): the forward pass also uses safe_x, so stable_sqrt(0) = 1
@jax.custom_jvp
def stable_sqrt_course(x):
is_zero = (x == 0.0)
safe_x = jnp.where(is_zero, 1.0, x)
return jnp.sqrt(safe_x)
@stable_sqrt_course.defjvp
def stable_sqrt_course_jvp(primals, tangents):
x, = primals
t, = tangents
primal_out = stable_sqrt_course(x)
is_zero = (x == 0.0)
safe_x = jnp.where(is_zero, 1.0, x)
deriv = 0.5 * jax.lax.rsqrt(safe_x)
tangent_out = jnp.where(is_zero, 0.0, deriv * t)
return primal_out, tangent_out
print("Original version, jvp at x=0:", jax.jvp(stable_sqrt_course, (0.0,), (1.0,)))
# ✅ Recommended: real sqrt in the forward pass; safe_x only in the derivative formula
@jax.custom_jvp
def stable_sqrt(x):
return jnp.sqrt(x)
@stable_sqrt.defjvp
def stable_sqrt_jvp(primals, tangents):
x, = primals
t, = tangents
is_zero = (x == 0.0)
safe_x = jnp.where(is_zero, 1.0, x)
tangent_out = jnp.where(is_zero, 0.0, 0.5 * lax.rsqrt(safe_x) * t) # by convention, the derivative at x=0 is 0
return jnp.sqrt(x), tangent_out
p0, t0 = jax.jvp(stable_sqrt, (0.0,), (1.0,))
p4, t4 = jax.jvp(stable_sqrt, (4.0,), (1.0,))
print("x=0:", p0, t0, " x=4:", p4, t4) # x=4: 0.5 * 4**-0.5 = 0.25
assert stable_sqrt_course(0.0) == 1.0 # the ❌ version's forward pass is wrong
assert p0 == 0.0 and t0 == 0.0
assert p4 == 2.0 and t4 == 0.25
# Compare: the built-in jnp.sqrt has an inf gradient at 0
assert jnp.isinf(jax.grad(jnp.sqrt)(0.0)) and jax.grad(stable_sqrt)(0.0) == 0.0
How it relates to custom_vjp:
⚠️ Note: it's easy to think "if you only use
jvp/jacfwd, definingcustom_jvpmay be enough; if you needjax.grad, you usually have to definecustom_vjp; JAX can sometimes transpose the JVP into a VJP automatically, but it's not always reliable, so defining both is safest". What actually happens:
- As long as the output of the JVP rule is linear in the tangent (a correct JVP rule always is), and the tangent side is computed only with JAX's built-in transposable primitives (the tangent isn't handed to a black box like
pure_callback), thecustom_jvpfunction can be used withgrad/vjp/jacrev. JAX derives the reverse rule automatically via the "linearize + transpose" from Section 4.1; it's not a "sometimes" thing (for counterexamples see the Section 4.7 extra "JVP rules must be linear in the tangent").- You can't define both on the same function: a
custom_jvpobject has nodefvjpmethod at all. A function gets one or the other.- Conversely,
custom_vjpfunctions don't support forward mode: callingjvp/jacfwd/hessianon them raisesTypeError.So the rule of thumb is: use
custom_jvpwhenever you can, and forward, reverse and higher-order derivatives all work. Only reach forcustom_vjpwhen you need to control the backward computation itself (implicit functions, gradient clipping, custom residuals to save memory, calling external code in the backward pass).
Functions with a custom JVP compose correctly with jit, vmap and so on. Their typical use cases are the same few: numerical stability, performance, handling discontinuities, and interfacing with external code.
# 1) custom_jvp gets reverse mode automatically (grad is obtained by transposition)
assert jax.grad(stable_sqrt)(4.0) == 0.25
assert jnp.allclose(jax.jacrev(stable_sqrt)(4.0), jax.jacfwd(stable_sqrt)(4.0))
# 2) A custom_jvp object has no defvjp: you can't define both
assert not hasattr(stable_sqrt, "defvjp")
print("Methods of a custom_jvp object:", [m for m in dir(stable_sqrt) if m.startswith("def")])
# 3) custom_vjp doesn't support forward mode: jvp / jacfwd / hessian all raise TypeError
n_err = 0
for name, call in [("jvp", lambda: jax.jvp(log1pexp, (1.0,), (1.0,))),
("jacfwd", lambda: jax.jacfwd(log1pexp)(1.0)),
("hessian", lambda: jax.hessian(log1pexp)(1.0))]:
try:
call()
except TypeError as e:
n_err += 1
print(f"{name:8s} -> TypeError: {str(e)[:70]}")
assert n_err == 3
🔍 Extra: a numerically stable softplus, and defjvps¶
⚪ Optional: when writing a
custom_jvpfor an element-wise function
The most common way to write it for element-wise functions is defjvps: provide one function jvp_i(t_i, ans, *primals) for each positional argument, returning that argument's tangent's contribution to the output tangent. JAX adds up the contributions from all arguments.
ansis the forward output, and you can reuse it in the rule to avoid recomputation (e.g. $\sigma'(x)=\sigma(1-\sigma)$ just usesans);- If an argument doesn't need a derivative, pass
Nonein its position (its gradient is 0); - The forward
primal_outis computed automatically by the original function, so the rule doesn't return it.
$$\mathrm{softplus}(x)=\log(1+e^x)=\max(x,0)+\log\big(1+e^{-|x|}\big),\qquad \mathrm{softplus}'(x)=\sigma(x)$$
@jax.custom_jvp
def softplus(x):
# Numerically stable forward pass: max(x,0) + log1p(exp(-|x|)), never overflows
return jnp.maximum(x, 0.0) + jnp.log1p(jnp.exp(-jnp.abs(x)))
softplus.defjvps(lambda t, ans, x: t * jax.nn.sigmoid(x))
xs = jnp.array([-100.0, -10.0, 0.0, 10.0, 100.0])
print("softplus:", softplus(xs))
print("grad :", jax.vmap(jax.grad(softplus))(xs))
assert jnp.allclose(softplus(xs), jax.nn.softplus(xs))
assert jnp.allclose(jax.vmap(jax.grad(softplus))(xs), jax.nn.sigmoid(xs))
assert jnp.all(jnp.isfinite(jax.vmap(jax.grad(jax.grad(softplus)))(xs))) # second derivatives work too
# Both forward and reverse work, and it passes check_grads (checking fwd and rev, second order)
check_grads(softplus, (jnp.float32(1.3),), order=2, modes=["fwd", "rev"])
# Reuse the forward result through ans: sigmoid' = ans * (1 - ans); with several arguments you can pass None for one
@jax.custom_jvp
def my_sigmoid(x):
return 1.0 / (1.0 + jnp.exp(-x))
my_sigmoid.defjvps(lambda t, ans, x: t * ans * (1.0 - ans))
@jax.custom_jvp
def mul_no_grad_y(x, y):
return x * y
mul_no_grad_y.defjvps(lambda tx, ans, x, y: tx * y, None) # no derivative for y
assert jnp.allclose(jax.grad(my_sigmoid)(0.3), jax.grad(jax.nn.sigmoid)(0.3))
assert jax.grad(mul_no_grad_y, argnums=(0, 1))(2.0, 3.0) == (3.0, 0.0)
print("defjvps checks passed")
🔍 Extra: JVP rules must be linear in the tangent¶
🟢 Must-read
Automatic transposition requires the rule's tangent_out to be a linear function of the tangent. Inside the rule you can do anything with the primal, including while_loop and external callbacks, as long as that computation doesn't depend on the tangent. If you write a rule that is nonlinear in the tangent (e.g. t**2), jvp still runs (though the result isn't mathematically a derivative), but grad raises NotImplementedError during transposition. Likewise, the tangent side must be expressed only with JAX primitives: if you hand the tangent to jax.pure_callback, jvp still runs but grad raises ValueError: Pure callbacks do not support transpose during transposition (verified below).
@jax.custom_jvp
def nonlinear_rule(x):
return jnp.sin(x)
@nonlinear_rule.defjvp
def _nonlinear_rule_jvp(primals, tangents):
x, = primals
t, = tangents
return jnp.sin(x), jnp.cos(x) * t**2 # ❗ nonlinear in the tangent
print("jvp runs:", jax.jvp(nonlinear_rule, (1.0,), (2.0,)))
try:
jax.grad(nonlinear_rule)(1.0)
raised = False
except NotImplementedError as e:
raised = True
print("grad -> NotImplementedError:", str(e)[:70])
assert raised
# Using while_loop on the primal inside the rule is perfectly fine: the tangent part is still linear
@jax.custom_jvp
def sin_with_loop(x):
return jnp.sin(x)
@sin_with_loop.defjvp
def _sin_with_loop_jvp(primals, tangents):
x, = primals
t, = tangents
d = lax.while_loop(lambda c: c[0] < 3, lambda c: (c[0] + 1, c[1]), (0, jnp.cos(x)))[1]
return jnp.sin(x), d * t
assert jnp.allclose(jax.grad(sin_with_loop)(1.0), jnp.cos(1.0))
# Handing the tangent to pure_callback: still linear mathematically, jvp runs, but grad fails at transposition
@jax.custom_jvp
def tangent_via_callback(x):
return jnp.sin(x)
@tangent_via_callback.defjvp
def _tangent_via_callback_jvp(primals, tangents):
x, = primals
t, = tangents
t_cb = jax.pure_callback(lambda a: np.asarray(a), jax.ShapeDtypeStruct(t.shape, t.dtype), t)
return jnp.sin(x), jnp.cos(x) * t_cb
x1 = jnp.float32(1.0)
assert jnp.allclose(jax.jvp(tangent_via_callback, (x1,), (x1,))[1], jnp.cos(1.0))
try:
jax.grad(tangent_via_callback)(x1)
raised = False
except ValueError as e:
raised = True
print("Tangent through pure_callback → grad:", str(e)[:60])
assert raised
4.8 Differentiating through control-flow primitives¶
JAX's autodiff is tightly integrated with lax.cond, lax.scan, lax.fori_loop and lax.while_loop. These primitives record the whole control structure into the computation graph during tracing, so you can differentiate through them under jit too. (For how Python's native if/for interact with grad, see Section 3.7 of Basics 03: Automatic differentiation with grad.)
lax.cond: JAX traces both true_fun and false_fun, and the derivative function contains the gradient logic for both branches, picking one at run time based on the predicate.
- Both branches must produce outputs with the same structure (same shapes and dtypes);
- Differentiating with respect to the predicate usually makes no sense: the predicate is a boolean, so the function is piecewise constant in the variable that decides it, and the gradient is 0.
Example: when $y>0$, $f=\sin(2x)$ with derivative $2\cos(2x)$, which is 0 at $x=\pi/4$; when $y\le 0$, $f=\cos(x/2)$ with derivative $-\tfrac12\sin(x/2)$, about $-0.191$ at $x=\pi/4$.
# Course code: differentiating lax.cond
def conditional_computation(x, y):
return lax.cond(y > 0,
lambda op: jnp.sin(op * 2.0), # true branch
lambda op: jnp.cos(op / 2.0), # false branch
x)
grad_x = jax.grad(conditional_computation, argnums=0)
grad_y = jax.grad(conditional_computation, argnums=1) # w.r.t. y, which decides the predicate: gradient is 0
x_val, y_val_pos, y_val_neg = jnp.pi / 4.0, 1.0, -1.0
print(f"f(x, y>0) = {conditional_computation(x_val, y_val_pos):.4f}, df/dx = {grad_x(x_val, y_val_pos):.4f}")
print(f"f(x, y<0) = {conditional_computation(x_val, y_val_neg):.4f}, df/dx = {grad_x(x_val, y_val_neg):.4f}")
print(f"df/dy (y>0) = {grad_y(x_val, y_val_pos)}, df/dy (y<0) = {grad_y(x_val, y_val_neg)}")
assert jnp.allclose(grad_x(x_val, y_val_pos), 0.0, atol=1e-6)
assert jnp.allclose(grad_x(x_val, y_val_neg), -0.5 * jnp.sin(jnp.pi / 8))
assert grad_y(x_val, y_val_pos) == 0.0 and grad_y(x_val, y_val_neg) == 0.0
lax.while_loop: it's tempting to compare it to BPTT in an RNN, assume the backward pass of a while_loop will "save every iteration's intermediates", and call jax.grad on a while_loop directly. For example, differentiate $\sum_{i=0}^{4} \text{scale}\cdot i$ and expect a gradient of $0+1+2+3+4=10$, wrapping it in jax.checkpoint if memory runs short.
⚠️ Note:
lax.while_loopdoes not support reverse-mode differentiation.jax.gradraisesValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop valuesright away. The reason is that the number of iterations is only known at run time, so JAX can't allocate buffers for the intermediates in advance.while_looponly supports forward mode (jvp/jacfwd). Wrapping the loop body injax.checkpointdoesn't help either; you get the same error. The right way to write this example:
- If the number of iterations is known at trace time: use
lax.fori_loopwith static Python integer bounds (it lowers toscaninternally), or uselax.scandirectly;- If the number of iterations depends on convergence: use
custom_vjpwith the implicit function theorem (Section 4.6), or stick to forward mode.
# ❌ Old style (outdated/wrong): reverse-mode differentiation of lax.while_loop
def loop_sum_param(scale, n_iters):
init_state = (0, 0.0)
def cond_fun(state):
i, _ = state
return i < n_iters
def body_fun(state):
i, current_sum = state
return i + 1, current_sum + scale * i
_, final_sum = lax.while_loop(cond_fun, body_fun, init_state)
return final_sum
grad_loop_sum = jax.grad(loop_sum_param, argnums=0)
try:
print(grad_loop_sum(2.0, 5))
except ValueError as e:
print("ValueError:", e)
# Wrapping in jax.checkpoint doesn't work either
try:
jax.grad(lambda s: lax.while_loop(lambda c: c[0] < 5,
jax.checkpoint(lambda c: (c[0] + 1, c[1] + s * c[0])),
(0, 0.0))[1])(2.0)
except ValueError as e:
print("Still a ValueError with checkpoint")
# ✅ Recommended
scale_val, iters = 2.0, 5
print("Forward value:", loop_sum_param(scale_val, iters)) # 2*(0+1+2+3+4) = 20
# (a) while_loop only supports forward mode: jvp / jacfwd work
_, tangent = jax.jvp(lambda s: loop_sum_param(s, iters), (scale_val,), (1.0,))
g_fwd = jax.jacfwd(loop_sum_param)(scale_val, iters)
# (b) fori_loop with static bounds (lowers to scan internally) supports reverse mode
def loop_sum_fori(scale, n_iters):
return lax.fori_loop(0, n_iters, lambda i, acc: acc + scale * i, 0.0)
g_rev = jax.grad(loop_sum_fori)(scale_val, iters) # n_iters is a Python int (static)
# (c) The equivalent scan version
def loop_sum_scan(scale, n_iters):
acc, _ = lax.scan(lambda acc, i: (acc + scale * i, None), 0.0, jnp.arange(n_iters))
return acc
g_scan = jax.grad(loop_sum_scan, argnums=0)(scale_val, iters)
print("jvp:", tangent, " jacfwd:", g_fwd, " grad(fori, static):", g_rev, " grad(scan):", g_scan)
assert loop_sum_param(scale_val, iters) == 20.0
assert tangent == 10.0 and g_fwd == 10.0 and g_rev == 10.0 and g_scan == 10.0
# Trap: inside jit, n_iters becomes a tracer -> fori_loop falls back to while_loop -> reverse mode fails again
try:
jax.jit(jax.grad(loop_sum_fori))(scale_val, iters)
raised = False
except ValueError:
raised = True
print("grad(fori_loop) raises under jit when n_iters is dynamic:", raised)
assert raised
# Fix: declare n_iters as a static argument
g_jit = jax.jit(jax.grad(loop_sum_fori), static_argnums=1)(scale_val, iters)
assert g_jit == 10.0
lax.scan: differentiation through scan is well defined and efficient, and JAX automatically handles gradient flow through the carry and the per-step outputs $y_t$. For fixed-length sequences scan is usually a better fit than while_loop (especially since the latter can't do reverse mode); for long sequences you can apply jax.checkpoint to the step function to save memory (the extra below only gives the key points; the full treatment is in Advanced Chapter 6 §6.6).
The example is a minimal RNN: $h_t=\tanh(0.5\,h_{t-1}+1.5\,x_t)$, $y_t = 2h_t$, loss $\sum_t y_t$, differentiated with respect to the initial state $h_0$ and the input sequence $x_{1:T}$.
⚠️ Note: the function that differentiates with respect to $h_0$ is named
grad_rnn_h0here, notgrad_rnn_params. It computes the gradient with respect to the initial state, not a parameter gradient: the weights in the step function are fixed constants, and there are no trainable parameters at all.
# Course code: differentiating lax.scan
def simple_rnn_step(carry, x_t):
prev_h = carry
weight_hh, weight_xh = 0.5, 1.5 # fixed parameters (not trainable)
new_h = jnp.tanh(prev_h * weight_hh + x_t * weight_xh)
y_t = new_h * 2.0
return new_h, y_t
def run_rnn(initial_state, inputs):
_, outputs_y = lax.scan(simple_rnn_step, initial_state, inputs)
return jnp.sum(outputs_y)
grad_rnn_h0 = jax.grad(run_rnn, argnums=0) # w.r.t. the initial state
grad_rnn_inputs = jax.grad(run_rnn, argnums=1) # w.r.t. the input sequence
h0 = jnp.zeros(())
xts = jnp.array([0.1, 0.2, -0.1, 0.3])
total_output, dh0, dxts = run_rnn(h0, xts), grad_rnn_h0(h0, xts), grad_rnn_inputs(h0, xts)
print("total_output:", total_output, " dL/dh0:", dh0, " dL/dx:", dxts)
# Cross-check against hand-written BPTT in NumPy float64
def bptt_np(h0, xs):
hs = [float(h0)]
for x in xs:
hs.append(np.tanh(0.5 * hs[-1] + 1.5 * x))
dh, dxs = 0.0, np.zeros(len(xs))
for t in reversed(range(len(xs))):
dh = dh + 2.0 # from y_t = 2 h_t
dpre = dh * (1 - hs[t + 1] ** 2)
dxs[t] = dpre * 1.5
dh = dpre * 0.5 # passed on to h_{t-1}
return 2 * sum(hs[1:]), dh, dxs
ref_out, ref_dh0, ref_dx = bptt_np(h0, np.asarray(xts, dtype=np.float64))
assert np.allclose(total_output, ref_out, rtol=1e-5)
assert np.allclose(dh0, ref_dh0, rtol=1e-5)
assert np.allclose(dxts, ref_dx, rtol=1e-5)
🔍 Extra: differentiability of each control-flow primitive at a glance (tested)¶
🟢 Must-read
| Primitive | Forward mode (jvp) | Reverse mode (grad) | Notes |
|---|---|---|---|
lax.cond / lax.switch | ✅ | ✅ | The predicate / index itself is boolean or integer, so you can't grad with respect to it (TypeError); differentiating with respect to a continuous variable that decides the predicate gives 0 |
lax.scan | ✅ | ✅ | Saves residuals for every step by default, memory $O(T)$ |
lax.fori_loop (static int bounds) | ✅ | ✅ | Lowers to scan internally |
lax.fori_loop (tracer bounds) | ✅ | ❌ ValueError | Lowers to while_loop internally |
lax.while_loop | ✅ | ❌ ValueError | Unknown number of iterations; use scan with an upper bound, or custom_vjp + the implicit function theorem |
🔍 Extra: jax.checkpoint + scan to save backward-pass memory¶
⚪ Optional: when backprop through a long sequence or a deep
scanruns out of memory
By default, the backward pass of scan stores the residuals of every step, so memory is $O(T\cdot\text{per-step intermediates})$. Once you wrap the step function in jax.checkpoint (alias jax.remat), the forward pass only saves each step's carry, and the backward pass recomputes the in-step intermediates, trading roughly one extra forward pass for a lot of memory. A quantitative comparison with memory_analysis() and print_saved_residuals, rematerialization policies (policy), and the standard "scan over stacked layers + checkpoint" pattern are all collected in Advanced Chapter 6 §6.6.
4.9 Handling non-differentiable functions¶
Autodiff works by recursively applying the chain rule to each primitive, which assumes every primitive has a JVP/VJP rule. Some operations are mathematically non-differentiable, some act on integers, and sometimes we deliberately cut the gradient for modeling or performance reasons.
Why differentiation can fail falls into three categories:
- Mathematical discontinuities: rounding (
round,floor,ceil);signat 0; booleans produced by comparisons and then used incondor arithmetic;argmax/argminreturn integer indices, which can jump when the input changes slightly. - Integer operations: type conversions and indexing by integer value have no derivative in the usual sense.
- The primitive has no differentiation rule: differentiating it raises an error.
JAX's default behavior: for operations like round and integer conversion whose derivative is "zero almost everywhere", JAX defines the gradient to be 0, so jax.grad(lambda x: jnp.round(x))(x) gives 0.0.
⚠️ Note: when an operation has no differentiation rule, the error isn't always a
TypeError; it depends on the case:
gradwith respect to an integer input:TypeError(suggestsallow_int=Trueorvjp);gradof an integer output (e.g. the index returned directly byargmax):TypeError;jax.pure_callbackhas no JVP:ValueError(suggestscustom_jvp);- A missing transpose rule (a nonlinear JVP rule, or calling
linear_transposeon a nonlinear function):NotImplementedError;- Reverse mode through
while_loop:ValueError.
errors = {}
def record(name, fn):
try:
fn()
errors[name] = "no error"
except Exception as e:
errors[name] = type(e).__name__
record("grad int input", lambda: jax.grad(lambda x: x * 2)(3))
record("grad int output (argmax)", lambda: jax.grad(lambda x: jnp.argmax(x))(jnp.array([0.3, 2.0, 1.0])))
record("grad pure_callback", lambda: jax.grad(
lambda x: jax.pure_callback(lambda v: np.sin(v).astype(np.float32),
jax.ShapeDtypeStruct((), jnp.float32), x))(1.0))
record("grad nonlinear JVP rule", lambda: jax.grad(nonlinear_rule)(1.0))
record("grad while_loop", lambda: jax.grad(loop_sum_param)(2.0, 5))
for k, v in errors.items():
print(f"{k:24s} -> {v}")
assert errors == {"grad int input": "TypeError", "grad int output (argmax)": "TypeError",
"grad pure_callback": "ValueError", "grad nonlinear JVP rule": "NotImplementedError",
"grad while_loop": "ValueError"}
🔍 Extra: tested gradient behavior of common "non-differentiable" ops¶
🟢 Must-read
| Operation | Gradient behavior (tested on JAX 0.12) |
|---|---|
round / floor / ceil / sign | 0 everywhere (including at the jumps) |
argmax (after casting to float) | 0; calling grad directly on the integer output raises TypeError |
sort | Nonzero: the gradient flows back to the original positions through the sort permutation (sort is piecewise linear in its input) |
lax.top_k | The values part works like sort: the gradient is scattered back to the selected positions, and unselected positions get 0; the indices part is integer and has no gradient |
max (with tied maxima) | Split evenly across tied positions: [1,2,2] → [0, 0.5, 0.5] |
abs at 0 | 1.0 (picks +1 from the subgradient) |
relu at 0 | 0.0 |
xx = jnp.array([0.3, 2.0, 1.0])
w = jnp.array([1.0, 2.0, 3.0])
results = {
"round": jax.grad(lambda x: jnp.round(x).sum())(xx),
"floor": jax.grad(lambda x: jnp.floor(x).sum())(xx),
"sign": jax.grad(lambda x: jnp.sign(x).sum())(xx),
"argmax(float)": jax.grad(lambda x: jnp.argmax(x).astype(jnp.float32))(xx),
"sort·w": jax.grad(lambda x: (jnp.sort(x) * w).sum())(xx),
"top_k(2)·[10,1]": jax.grad(lambda x: (lax.top_k(x, 2)[0] * jnp.array([10.0, 1.0])).sum())(xx),
"max with ties": jax.grad(jnp.max)(jnp.array([1.0, 2.0, 2.0])),
}
for k, v in results.items():
print(f"{k:16s} {v}")
print("abs'(0) =", jax.grad(jnp.abs)(0.0), " relu'(0) =", jax.grad(jax.nn.relu)(0.0))
for k in ["round", "floor", "sign", "argmax(float)"]:
assert jnp.all(results[k] == 0.0)
# After sorting the order is [0.3, 1.0, 2.0], i.e. original positions [0, 2, 1], so the gradient at the original positions is [w0, w2, w1] = [1, 3, 2]
assert jnp.allclose(results["sort·w"], jnp.array([1.0, 3.0, 2.0]))
# top_2 selects 2.0 (position 1, weight 10) and 1.0 (position 2, weight 1)
assert jnp.allclose(results["top_k(2)·[10,1]"], jnp.array([0.0, 10.0, 1.0]))
assert jnp.allclose(results["max with ties"], jnp.array([0.0, 0.5, 0.5]))
assert jax.grad(jnp.abs)(0.0) == 1.0 and jax.grad(jax.nn.relu)(0.0) == 0.0
Stopping gradients explicitly with jax.lax.stop_gradient: the forward pass is the identity, and the gradient in the backward pass (and in JVPs) is 0. In other words, as far as differentiation is concerned, whatever it wraps is a constant. (For the basics, see Section 3.8 of Basics 03: Automatic differentiation with grad.)
Example: $f(x)=\sin x\cos x$. The normal gradient is $\cos^2x-\sin^2x=\cos 2x$, which is 0 at $x=\pi/4$. Once $\cos x$ is wrapped in stop_gradient, the gradient is only $\cos x\cdot\cos x = 0.5$.
When to stop gradients, five common cases:
- It's mathematically required (e.g. round);
- Modeling assumptions (e.g. target networks in reinforcement learning);
- Truncating gradient flow through unrolled loops (truncated BPTT);
- Performance: a branch cut off by
stop_gradientis treated as a constant during differentiation, so JAX doesn't generate any backward computation for it at all (rather than computing it and multiplying by 0); - Numerical stability: isolating the parts that would produce NaN/Inf.
# Course code: stop_gradient
def f_normal(x):
y = jnp.sin(x)
z = jnp.cos(x)
return y * z # gradient: cos^2 - sin^2 = cos(2x)
def f_stopped(x):
y = jnp.sin(x)
z = jax.lax.stop_gradient(jnp.cos(x))
return y * z # gradient: cos(x) * z
x_val = jnp.pi / 4.0
print(f"x = {x_val:.3f}, f_normal = {f_normal(x_val):.4f}, f_stopped = {f_stopped(x_val):.4f} (same forward value)")
print(f"grad normal = {jax.grad(f_normal)(x_val):.4f}, grad stopped = {jax.grad(f_stopped)(x_val):.4f}")
assert jnp.allclose(f_normal(x_val), f_stopped(x_val))
assert jnp.allclose(jax.grad(f_normal)(x_val), 0.0, atol=1e-6)
assert jnp.allclose(jax.grad(f_stopped)(x_val), 0.5)
# Also 0 in a JVP
assert jax.jvp(lax.stop_gradient, (1.0,), (1.0,))[1] == 0.0
Alternatives for non-differentiable operations:
- Smooth approximations: e.g. approximate a step function with $\sigma(kx)$; the larger $k$, the closer it gets;
- Custom gradients, especially the Straight-Through Estimator (STE): use the non-differentiable op (e.g. rounding) in the forward pass, and pass the upstream gradient $g$ straight back in the backward pass, as if the op were the identity.
Below we implement STE with custom_vjp: ste_example(x) = round(x) * x. At $x=2.7$, the forward value is $3\times2.7=8.1$ and the STE gradient is $\mathrm{round}(x)+x=3+2.7=5.7$ (the true gradient is $\mathrm{round}(x) = 3$). fwd calls round_straight_through(x) itself directly. That's fine, because calling it inside fwd doesn't trigger the custom rule again.
# Course code: STE with custom_vjp
@jax.custom_vjp
def round_straight_through(x):
"""round in the forward pass, identity in the backward pass"""
return jnp.round(x)
def round_straight_through_fwd(x):
return round_straight_through(x), None # no residuals needed
def round_straight_through_bwd(residuals, g):
return (g,) # gradient passes straight through
round_straight_through.defvjp(round_straight_through_fwd, round_straight_through_bwd)
def ste_example(x):
rounded_x = round_straight_through(x)
return rounded_x * x
x_val = 2.7
print("ste_example(2.7) =", ste_example(x_val), " grad =", jax.grad(ste_example)(x_val))
assert jnp.allclose(ste_example(x_val), 8.1)
assert jnp.allclose(jax.grad(ste_example)(x_val), 5.7)
assert jnp.allclose(jax.grad(lambda x: jnp.round(x) * x)(x_val), 3.0) # without STE, round's gradient is 0
🔍 Extra: STE with stop_gradient, soft-argmax, and Gumbel-softmax¶
⚪ Optional: for quantization-aware training, discrete sampling, and other cases where you "specify the gradient by hand"
One-line STE: x + stop_gradient(round(x) - x). In the forward pass, $x + (\mathrm{round}(x)-x) = \mathrm{round}(x)$; in the backward pass the stop_gradient part is a constant, leaving just the identity gradient of $x$. It needs no custom_vjp, and it supports both forward and reverse mode. Quantization-aware training (QAT) and VQ-VAE often write it this way.
soft-argmax: $\mathrm{softargmax}_\tau(x)=\sum_i i\cdot\mathrm{softmax}(x/\tau)_i$. As $\tau\to0$ it approaches argmax, and it is differentiable everywhere.
Gumbel-softmax (the Concrete distribution): to sample differentiably from a categorical distribution $\pi$, add Gumbel noise $g_i\sim\mathrm{Gumbel}(0,1)$ and then apply a softmax with temperature: $y=\mathrm{softmax}\big((\log\pi+g)/\tau\big)$. With hard=True, the forward output is one-hot and the backward pass uses the soft version's gradient (another kind of STE). The Gumbel-max trick guarantees that $\arg\max_i(\log\pi_i+g_i)$ follows $\pi$ exactly.
# STE with stop_gradient: identical to the custom_vjp version, and it supports forward mode
def round_ste(x):
return x + lax.stop_gradient(jnp.round(x) - x)
assert round_ste(2.7) == 3.0
assert jnp.allclose(jax.grad(lambda x: round_ste(x) * x)(2.7), 5.7)
assert jax.jvp(round_ste, (2.7,), (1.0,))[1] == 1.0 # forward mode works too
# Smooth approximation of a step function: sigmoid(kx)
step_smooth = lambda x, k: jax.nn.sigmoid(k * x)
print("sigmoid(kx) at x=0.1:", [float(step_smooth(0.1, k)) for k in (1, 10, 100)])
# soft-argmax
def soft_argmax(x, tau):
p = jax.nn.softmax(x / tau)
return jnp.sum(p * jnp.arange(x.shape[0]))
x3 = jnp.array([0.3, 2.0, 1.0])
print("soft_argmax(τ=0.05) =", soft_argmax(x3, 0.05), " argmax =", jnp.argmax(x3))
print("grad soft_argmax(τ=0.5) =", jax.grad(soft_argmax)(x3, 0.5))
assert jnp.allclose(soft_argmax(x3, 0.05), 1.0, atol=1e-4)
assert jnp.any(jax.grad(soft_argmax)(x3, 0.5) != 0.0)
# Gumbel-softmax (hard version = one-hot forward + soft gradient)
def gumbel_softmax(key, logits, tau, hard=False):
g = jax.random.gumbel(key, logits.shape)
y = jax.nn.softmax((logits + g) / tau)
if hard:
y_hard = jax.nn.one_hot(jnp.argmax(y), logits.shape[-1])
y = y + lax.stop_gradient(y_hard - y)
return y
probs = jnp.array([0.2, 0.5, 0.3])
keys = jax.random.split(jax.random.key(0), 20000)
samples = jax.vmap(lambda k: gumbel_softmax(k, jnp.log(probs), 0.5, hard=True))(keys)
freq = samples.mean(axis=0)
print("hard Gumbel-softmax sample frequencies:", freq, " target:", probs)
assert jnp.allclose(samples.sum(axis=1), 1.0) and jnp.all((samples == 0) | (samples == 1)) # forward is one-hot
assert jnp.allclose(freq, probs, atol=0.015) # correct distribution
g_logits = jax.grad(lambda l: gumbel_softmax(jax.random.key(0), l, 0.5, hard=True)[1])(jnp.zeros(3))
print("Gradient w.r.t. logits (from the soft version):", g_logits)
assert jnp.any(g_logits != 0.0)
4.10 Practice: implementing custom gradients¶
This section applies custom_vjp to softplus. custom_vjp has four kinds of uses: numerical stability (autodiff gradients can overflow or underflow), performance (you already know a more efficient gradient algorithm), non-JAX code (you must supply the gradient by hand), and stopped or approximate gradients.
$$\mathrm{softplus}(x)=\log(1+e^x),\qquad \frac{d}{dx}\mathrm{softplus}(x)=\frac{e^x}{1+e^x}=\frac{1}{1+e^{-x}}=\sigma(x)$$
$\sigma(x)$ is numerically stable on its own, so it can be used directly in the custom VJP. The three parts:
- Decorate the original function with
@jax.custom_vjp; fwd: same arguments as the original function, returns(output, residuals);bwd: receivesresidualsand the upstream gradient $g = \partial\mathcal L/\partial y$, and returns a tuple of gradients for each positional argument, i.e. $\partial\mathcal L/\partial x = g\cdot \partial y/\partial x$.
⚠️ Note: don't expect JAX and XLA to "smartly" handle the overflow of $e^x$ for large $x$ so that default autodiff of $\log(1+e^x)$ gives a stable gradient. In float32 on CPU: the naive implementation
jnp.log(1 + jnp.exp(x))at $x=100$ givesinfin the forward pass andnanas the gradient. Even sneakier is $x=88$:exp(88)itself doesn't overflow yet, but the factor $\frac{1}{1+e^x}\approx 6\times10^{-39}$ in the gradient $e^x\cdot\frac{1}{1+e^x}$ is a subnormal number that the CPU flushes to 0, so the gradient silently becomes 0.0 (the true value is about 1). XLA does not rewrite it into a stable form.⚠️ Note: catching overflow with
except OverflowError: naive_grad = float('inf')and expectinginfdoesn't work either. JAX raises no exception at all, so that branch is dead code, and what you actually get isnan.
# Course code 1: naive implementation
def softplus_naive(x):
return jnp.log(1 + jnp.exp(x)) # unstable for large x
x_val = 10.0
print(f"softplus_naive({x_val}) = {softplus_naive(x_val)}")
print(f"Gradient at {x_val} (JAX autodiff): {jax.grad(softplus_naive)(x_val)}")
assert jnp.allclose(jax.grad(softplus_naive)(x_val), jax.nn.sigmoid(10.0))
# ❌ Old style (outdated/wrong): the forward pass is still naive; overflow is "caught" with except OverflowError (never triggers)
@jax.custom_vjp
def softplus_custom(x):
"""Softplus implementation for a custom VJP"""
return jnp.log(1 + jnp.exp(x)) # the forward pass is still the naive version (for demonstration only)
def softplus_fwd(x):
y = softplus_custom(x)
residuals = x # the backward pass needs x to compute sigmoid(x)
return y, residuals
def softplus_bwd(residuals, g):
x = residuals
grad_x = g * jax.nn.sigmoid(x) # stable
return (grad_x,)
softplus_custom.defvjp(softplus_fwd, softplus_bwd)
grad_softplus_custom = jax.grad(softplus_custom)
test_values = jnp.array([-10.0, 0.0, 10.0, 80.0, 88.0, 100.0]) # 88 added to show the gradient silently going wrong
print("Input | Naive grad (autodiff) | Custom grad (VJP) | Forward softplus_custom")
print("-" * 72)
overflow_caught = False
naive_grads = []
for x_val in test_values:
try:
naive_grad = jax.grad(softplus_naive)(x_val)
except OverflowError:
overflow_caught = True
naive_grad = float("inf")
naive_grads.append(naive_grad)
custom_grad = grad_softplus_custom(x_val)
print(f"{x_val:<5.1f} | {naive_grad:<21.8f} | {custom_grad:<18.8f} | {softplus_custom(x_val)}")
jit_grad_softplus_custom = jax.jit(jax.grad(softplus_custom))
print(f"\nJIT-compiled custom gradient at 10.0: {jit_grad_softplus_custom(10.0):.8f}")
print("Did except OverflowError trigger:", overflow_caught)
# ✅ Recommended: stable forward pass (logaddexp) + custom VJP (sigmoid)
@jax.custom_vjp
def softplus_stable(x):
return jnp.logaddexp(x, 0.0)
def softplus_stable_fwd(x):
return softplus_stable(x), x
def softplus_stable_bwd(x, g):
return (g * jax.nn.sigmoid(x),)
softplus_stable.defvjp(softplus_stable_fwd, softplus_stable_bwd)
g_stable = jax.vmap(jax.grad(softplus_stable))(test_values)
print("Stable version, forward:", softplus_stable(test_values))
print("Stable version, gradient:", g_stable)
naive_grads = jnp.array(naive_grads)
assert not overflow_caught # the except branch never ran
assert jnp.isnan(naive_grads[-1]) # x=100: nan, not inf
assert abs(float(naive_grads[-2]) - 1.0) > 0.5 # at x=88 the gradient is silently wrong (0.0 on this machine)
assert jnp.isinf(softplus_custom(100.0)) # this version's forward pass overflows at x=100
assert jnp.allclose(g_stable, jax.nn.sigmoid(test_values)) # the stable version's gradient is correct everywhere
assert jnp.allclose(softplus_stable(test_values), jax.nn.softplus(test_values))
assert jnp.allclose(jax.jit(jax.grad(softplus_stable))(10.0), 0.9999546, atol=1e-6)
check_grads(softplus_stable, (jnp.float32(2.0),), order=1, modes=["rev"])
Verification and analysis: in the normal range, the custom gradient matches the automatic one (to floating-point precision); the backward pass explicitly uses jax.nn.sigmoid, which keeps it stable.
⚠️ Note:
softplus_customonly changes the gradient; its forward pass is still the naive implementation, which givesinfat $x=100$. Production code needs a stable forward pass: the ✅ version above is a complete implementation, with a stable forward (logaddexp) plus a custom VJP, and it passescheck_grads. You can also just usejax.nn.softplus.
The whole mechanism: softplus_fwd caches x, and softplus_bwd computes g * sigmoid(x).
More things to consider:
- Residuals should hold only what the backward pass needs: storing large intermediate arrays costs extra memory, and sometimes recomputing them in
bwdis cheaper. (For example, softplus only storesx, notexp(x).) - Multiple inputs / outputs:
fwdreceives all inputs; the tuple returned bybwdmust have as many elements as there are positional arguments;ghas the same structure as the original function's output. jax.custom_jvp: computes(y, y_dot)from(x, x_dot); see Section 4.7.- Always test custom rules thoroughly against JAX's autodiff or numerical gradients.
⚠️ Note: for outputs that aren't on the gradient path, the corresponding
gthatbwdreceives is notNone. By default (symbolic_zeros=False) it's an array of zeros. Only withdefvjp(fwd, bwd, symbolic_zeros=True)do you get symbolic zeros, which you can then use to skip computation.
# For an output that doesn't contribute to the loss, bwd receives a zero array, not None
seen = {}
@jax.custom_vjp
def two_outputs(x):
return jnp.sin(x), jnp.cos(x)
def two_outputs_fwd(x):
return two_outputs(x), x
def two_outputs_bwd(x, g):
g_sin, g_cos = g
seen["types"] = (type(g_sin).__name__, type(g_cos).__name__)
seen["g_cos"] = g_cos
return (g_sin * jnp.cos(x) - g_cos * jnp.sin(x),)
two_outputs.defvjp(two_outputs_fwd, two_outputs_bwd)
gx = jax.grad(lambda x: two_outputs(x)[0])(1.0) # use only the first output
print("Types of g received by bwd:", seen["types"], " g_cos =", seen["g_cos"])
assert seen["g_cos"] is not None and seen["g_cos"] == 0.0
assert jnp.allclose(gx, jnp.cos(1.0))
🔍 Putting it together: an implicit layer (Deep Equilibrium, DEQ) with custom_vjp¶
The output of an implicit layer is defined as a fixed point instead of a finite number of computation steps:
$$z^\star = \tanh(Wz^\star + Ux),\qquad \text{loss} = \|z^\star\|^2$$
- Forward: iterate with
while_loopuntil convergence (the number of iterations depends on the data); - Backward (implicit function theorem, same as Section 4.6): let $T(\theta, x, z)=\tanh(Wz+Ux)$; first solve the adjoint fixed point $u = g + (\partial T/\partial z)^\top u$, then compute $\bar\theta = u^\top\partial T/\partial\theta$ and $\bar x = u^\top \partial T/\partial x$;
- Memory: only $z^\star$ is saved, independent of the number of iterations; the unrolled version has to save every step.
When $\|W\|$ is small enough, $T$ is a contraction in $z$, so both the forward and adjoint iterations converge. Parameters are passed in as function arguments (not captured in a closure), and the result is cross-checked against "unroll 100 steps and backprop directly".
def deq_cell(params, x, z):
W, U = params
return jnp.tanh(W @ z + U @ x)
@partial(jax.custom_vjp, nondiff_argnums=(0,))
def deq_solve(T, params, x):
W, _ = params
def cond(c):
z, z_prev, i = c
return (jnp.max(jnp.abs(z - z_prev)) > 1e-6) & (i < 500)
def body(c):
z, _, i = c
return T(params, x, z), z, i + 1
z0 = jnp.zeros(W.shape[0], W.dtype)
z_star, _, n_iter = lax.while_loop(cond, body, (T(params, x, z0), z0, 0))
return z_star
def deq_fwd(T, params, x):
z_star = deq_solve(T, params, x)
return z_star, (params, x, z_star)
def deq_bwd(T, res, g):
params, x, z_star = res
_, vjp_z = jax.vjp(lambda z: T(params, x, z), z_star)
def cond(c):
u, u_prev, i = c
return (jnp.max(jnp.abs(u - u_prev)) > 1e-7) & (i < 500)
def body(c):
u, _, i = c
return g + vjp_z(u)[0], u, i + 1
u = lax.while_loop(cond, body, (g + vjp_z(g)[0], g, 0))[0]
_, vjp_px = jax.vjp(lambda p, xx: T(p, xx, z_star), params, x)
return vjp_px(u) # (gradient of params, gradient of x)
deq_solve.defvjp(deq_fwd, deq_bwd)
def loss_implicit(params, x):
return jnp.sum(deq_solve(deq_cell, params, x) ** 2)
def loss_unrolled(params, x, n_steps=100):
z = lax.fori_loop(0, n_steps, lambda i, z: deq_cell(params, x, z), jnp.zeros(params[0].shape[0]))
return jnp.sum(z ** 2)
k1, k2, k3 = jax.random.split(jax.random.key(42), 3)
d, m = 32, 8
W = jax.random.normal(k1, (d, d)) * 0.4 / jnp.sqrt(d) # small spectral norm guarantees a contraction
U = jax.random.normal(k2, (d, m)) / jnp.sqrt(m)
x_in = jax.random.normal(k3, (m,))
params = (W, U)
(gW_i, gU_i), gx_i = jax.jit(jax.grad(loss_implicit, argnums=(0, 1)))(params, x_in)
(gW_u, gU_u), gx_u = jax.jit(jax.grad(loss_unrolled, argnums=(0, 1)))(params, x_in)
for name, a, b in [("dW", gW_i, gW_u), ("dU", gU_i, gU_u), ("dx", gx_i, gx_u)]:
print(f"{name}: max abs difference {float(jnp.max(jnp.abs(a - b))):.2e}, magnitude {float(jnp.max(jnp.abs(b))):.3f}")
assert jnp.allclose(a, b, atol=1e-5, rtol=1e-4)
assert jnp.allclose(loss_implicit(params, x_in), loss_unrolled(params, x_in), atol=1e-6)
# Independent cross-check in NumPy float64: use the analytic implicit-function-theorem formula directly (explicit inverse, fine for small dimensions)
Wn, Un, xn = (np.asarray(a, dtype=np.float64) for a in (W, U, x_in))
z = np.zeros(d)
for _ in range(500):
z = np.tanh(Wn @ z + Un @ xn)
D = np.diag(1 - z**2) # tanh' at z*
g = 2 * z # dL/dz*
u = np.linalg.solve((np.eye(d) - D @ Wn).T, g) # u = (I - ∂T/∂z)^{-T} g
s = D @ u # passed back to the input of tanh
gW_ref, gU_ref, gx_ref = np.outer(s, z), np.outer(s, xn), Un.T @ s
assert np.allclose(gW_i, gW_ref, atol=1e-5)
assert np.allclose(gU_i, gU_ref, atol=1e-5)
assert np.allclose(gx_i, gx_ref, atol=1e-5)
print("Implicit-layer gradients match the NumPy float64 analytic solution")
# Saved residuals: the implicit version stores only z*, the unrolled version stores every step
print("\nResiduals saved by the implicit version:")
print_saved_residuals(loss_implicit, params, x_in)
Exercises¶
Section 4.10 implies 3 tasks, which are Course exercises 1–3 below; after those come 4 extra exercises. Each one has a reference solution and a self-check with assert.
Course exercise 1: compare the naive gradient with the custom gradient¶
In the normal range ($|x|\le 20$), verify that the gradient of softplus_custom from Section 4.10 matches the naive autodiff gradient.
Self-check interface: g_naive, g_custom: the gradients of the naive softplus_naive and the custom softplus_custom at 41 evenly spaced points on $[-20, 20]$, both with shape (41,).
# ✍️ Your answer: Course exercise 1
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Course exercise 1
try:
g_naive, g_custom
except NameError as e:
print(f"⏳ Course exercise 1: not answered yet: {e}")
else:
_xs = np.linspace(-20.0, 20.0, 41)
_sig = 1.0 / (1.0 + np.exp(-_xs)) # analytic derivative of softplus, σ(x), in float64
assert np.shape(g_naive) == (41,) and np.shape(g_custom) == (41,), "Both gradient arrays should have shape (41,), matching linspace(-20, 20, 41)"
assert np.allclose(g_naive, _sig, atol=1e-6), "g_naive doesn't match the analytic derivative σ(x): make sure you take grad of softplus_naive point by point"
assert np.allclose(g_custom, _sig, atol=1e-6), "g_custom doesn't match the analytic derivative σ(x): make sure you take grad of softplus_custom point by point"
print("✅ Course exercise 1: self-check passed")
Course exercise 2: choosing residuals — save or recompute¶
Write two custom_vjp versions of $f(x)=\sum_i \sin(x_i)^2$: A saves $\cos x$ and $\sin x$ as residuals; B saves only $x$ and recomputes in bwd. Compare how much data each saves with print_saved_residuals, and verify that the gradients are the same.
Self-check interface: sin2_save(x), sin2_recompute(x): two jax.custom_vjp functions, both computing $\sum_i\sin(x_i)^2$ (x is a 1D array).
# ✍️ Your answer: Course exercise 2
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Course exercise 2
try:
sin2_save, sin2_recompute
except NameError as e:
print(f"⏳ Course exercise 2: not answered yet: {e}")
else:
import io, re, contextlib
def _resid_elems(f, *args):
"""Output of print_saved_residuals → total number of saved residual elements"""
buf = io.StringIO()
with contextlib.redirect_stdout(buf):
print_saved_residuals(f, *args)
tot = 0
for line in buf.getvalue().splitlines():
m = re.match(r"\s*\w+\[([\d,]*)\]", line)
if m:
tot += int(np.prod([int(d) for d in m.group(1).split(",") if d] or [1]))
return tot
_x = jax.random.uniform(jax.random.key(42), (257,), minval=-3.0, maxval=3.0)
_x64 = np.asarray(_x, np.float64)
for _name, _f in [("sin2_save", sin2_save), ("sin2_recompute", sin2_recompute)]:
assert isinstance(_f, jax.custom_vjp), f"{_name} should be defined with jax.custom_vjp"
assert np.isclose(float(_f(_x)), np.sum(np.sin(_x64) ** 2), rtol=1e-5), f"{_name}: wrong forward value"
assert np.allclose(jax.jit(jax.grad(_f))(_x), np.sin(2 * _x64), atol=1e-5), f"{_name}: gradient doesn't match the analytic 2 sin x cos x"
_n_save, _n_re = _resid_elems(sin2_save, _x), _resid_elems(sin2_recompute, _x)
assert _n_save >= 2 * 257, f"sin2_save should save two arrays, sin x and cos x (it actually saved only {_n_save} elements)"
assert _n_re <= 257, f"sin2_recompute should save only x and recompute in bwd (it actually saved {_n_re} elements)"
print(f"Residual elements: A = {_n_save}, B = {_n_re}")
print("✅ Course exercise 2: self-check passed")
Course exercise 3: softplus for production¶
Implement a softplus whose forward pass doesn't overflow on $x\in[-1000,1000]$, whose gradient is given by custom_vjp, and which works with jit + vmap. (This is softplus_stable from Section 4.10; here we verify it on an extreme range.)
Self-check interface: softplus_prod(x): your implementation (scalar input, scalar output, gradient given by jax.custom_vjp).
# ✍️ Your answer: Course exercise 3
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Course exercise 3
try:
softplus_prod
except NameError as e:
print(f"⏳ Course exercise 3: not answered yet: {e}")
else:
assert isinstance(softplus_prod, jax.custom_vjp), "The exercise requires the gradient to be given by jax.custom_vjp"
_x = jnp.array([-1000.0, -300.0, -30.0, -1.0, 0.0, 0.5, 30.0, 300.0, 1000.0])
_x64 = np.asarray(_x, np.float64)
_y = jax.jit(jax.vmap(softplus_prod))(_x)
_g = jax.jit(jax.vmap(jax.grad(softplus_prod)))(_x)
assert np.all(np.isfinite(_y)), f"inf/nan in the forward pass: {_y}"
assert np.all(np.isfinite(_g)), f"inf/nan in the gradient: {_g}"
assert np.allclose(_y, np.logaddexp(_x64, 0.0), rtol=1e-6, atol=1e-6), "Forward values don't match the float64 reference for log(1+eˣ)"
assert np.allclose(_g, 0.5 * (1 + np.tanh(_x64 / 2)), atol=1e-6), "Gradient doesn't match σ(x)"
print("✅ Course exercise 3: self-check passed")
Extra exercise A: global-norm gradient clipping for a pytree¶
Implement clip_grad_global_norm(params, max_norm) with custom_vjp: identity in the forward pass; in the backward pass, if the global L2 norm of the whole pytree cotangent exceeds max_norm, scale it down proportionally to max_norm. This is a "per-layer insertion" style of clipping, with the same effect as optax clipping inside the optimizer. Verify with a dict of parameters: after clipping, the gradient's norm equals max_norm and its direction is unchanged.
Self-check interface: clip_grad_global_norm(params, max_norm): returns params (any pytree) unchanged in the forward pass, and clips the cotangent by global norm in the backward pass; max_norm is a Python float.
# ✍️ Your answer: Extra exercise A
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise A
try:
clip_grad_global_norm
except NameError as e:
print(f"⏳ Extra exercise A: not answered yet: {e}")
else:
_k = jax.random.split(jax.random.key(7), 6)
_p = {"a": jax.random.normal(_k[0], (4, 3)), "b": [jax.random.normal(_k[1], (5,)), jnp.array(2.0)]}
_c = {"a": jax.random.normal(_k[2], (4, 3)), "b": [jax.random.normal(_k[3], (5,)), jnp.array(-1.5)]}
_lin = lambda q: sum(jnp.vdot(u, v) for u, v in zip(jax.tree.leaves(_c), jax.tree.leaves(q))) # its gradient is exactly _c
_nc = float(np.sqrt(sum(np.sum(np.asarray(l) ** 2) for l in jax.tree.leaves(_c))))
_out = clip_grad_global_norm(_p, 1.0)
assert jax.tree.structure(_out) == jax.tree.structure(_p) and all(
np.array_equal(u, v) for u, v in zip(jax.tree.leaves(_out), jax.tree.leaves(_p))), "The forward pass should return params unchanged"
for _mn in [_nc / 4, _nc * 3]:
_g = jax.jit(jax.grad(lambda q: _lin(clip_grad_global_norm(q, _mn))))(_p)
_scale = min(1.0, _mn / _nc)
for u, v in zip(jax.tree.leaves(_g), jax.tree.leaves(_c)):
assert np.allclose(u, _scale * np.asarray(v), rtol=1e-5, atol=1e-6), (
f"With max_norm={_mn:.3f} (cotangent norm {_nc:.3f}) the gradient should be cotangent × {_scale:.3f}: "
+ ("above the threshold it should be scaled down proportionally" if _scale < 1 else "below the threshold the gradient should not change"))
print("✅ Extra exercise A: self-check passed")
Extra exercise B: coloring a pentadiagonal Jacobian¶
The function $f_i = x_i^3 + x_{i-2}x_{i+2} - x_{i-1}$ (out-of-range terms count as 0) has a pentadiagonal Jacobian (bandwidth 2). Use 5 colors and vmap(jvp) to compute the Jacobian for $n=50$, and cross-check against jacfwd. Think about it: with bandwidth $b$, what is the minimum number of colors? (Answer: $2b+1$.)
Self-check interface: f_penta(x): the pentadiagonal function above; colored_jacobian_banded(f, x, bandwidth): for an f with bandwidth bandwidth, compute the full $n\times n$ Jacobian with $2b+1$ colors and vmap(jvp).
# ✍️ Your answer: Extra exercise B
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise B
try:
f_penta, colored_jacobian_banded
except NameError as e:
print(f"⏳ Extra exercise B: not answered yet: {e}")
else:
_x = jax.random.normal(jax.random.key(3), (37,))
_xn = np.pad(np.asarray(_x, np.float64), 2)
_ref = _xn[2:-2] ** 3 + _xn[:-4] * _xn[4:] - _xn[1:-3] # f_i = x_i³ + x_{i-2}x_{i+2} − x_{i-1}
assert np.allclose(f_penta(_x), _ref, atol=1e-5), "f_penta's values don't match the definition (out-of-range terms should count as 0)"
_J = jax.jacfwd(f_penta)(_x)
assert np.allclose(colored_jacobian_banded(f_penta, _x, 2), _J, atol=1e-5), "Pentadiagonal case (bandwidth=2): Jacobian doesn't match jacfwd"
f_tri = lambda v: jnp.sin(v) + jnp.pad(v, 1)[:-2] * jnp.pad(v, 1)[2:] # tridiagonal
assert np.allclose(colored_jacobian_banded(f_tri, _x, 1), jax.jacfwd(f_tri)(_x), atol=1e-5), "Wrong Jacobian for the tridiagonal function (bandwidth=1)"
# With real color compression, too few colors make neighboring columns add onto each other → values inside the band go wrong too;
# if you just run jacfwd and cut out the band, this would instead come out "right".
_band1 = np.where(np.abs(np.subtract.outer(np.arange(37), np.arange(37))) <= 1, np.asarray(_J), 0.0)
assert not np.allclose(colored_jacobian_banded(f_penta, _x, 1), _band1, atol=1e-5), (
"With too few colors (only 3 for a pentadiagonal function) the result still equals the band of the true Jacobian: no real color compression is happening")
print("✅ Extra exercise B: self-check passed")
Extra exercise C: differentiating an ODE with the adjoint method¶
For $\dot z = \tanh(Wz)$, $z(0)=z_0$, integrate to $T=1$ with fixed-step RK4, with loss $L=\|z(T)\|^2$. Implement the continuous adjoint method with custom_vjp:
$$a(T)=\frac{\partial L}{\partial z(T)},\quad \dot a = -a^\top\frac{\partial f}{\partial z},\quad \frac{dL}{dW} = \int_0^T a^\top\frac{\partial f}{\partial W}\,dt$$
Integrate the augmented state $(z, a, \bar W)$ backward from $T$ to 0 ($z$ is integrated backward too, reconstructing it so the trajectory needn't be saved), giving $\partial L/\partial z_0 = a(0)$ and $\partial L/\partial W$. Then cross-check against "differentiate the unrolled RK4 version directly". The continuous adjoint and discrete backprop differ only by an $O(h^4)$ discretization error, so with $N=100$ steps the two should agree within $10^{-4}$.
Self-check interface: odeint_unrolled(z0, W): the unrolled RK4 version (global N_STEPS=100, T_END=1.0); odeint_adjoint(z0, W): the same forward pass, with a continuous-adjoint backward pass implemented via jax.custom_vjp. Both return $z(T)$; z0 has shape (d,) and W has shape (d, d).
# ✍️ Your answer: Extra exercise C
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise C
try:
odeint_unrolled, odeint_adjoint
except NameError as e:
print(f"⏳ Extra exercise C: not answered yet: {e}")
else:
from jax.experimental.ode import odeint as _dopri
import io, re, contextlib
def _resid_elems(f, *args):
"""Output of print_saved_residuals → total number of saved residual elements"""
buf = io.StringIO()
with contextlib.redirect_stdout(buf):
print_saved_residuals(f, *args)
tot = 0
for line in buf.getvalue().splitlines():
m = re.match(r"\s*\w+\[([\d,]*)\]", line)
if m:
tot += int(np.prod([int(d) for d in m.group(1).split(",") if d] or [1]))
return tot
assert isinstance(odeint_adjoint, jax.custom_vjp), "odeint_adjoint should be defined with jax.custom_vjp"
_k1, _k2 = jax.random.split(jax.random.key(123))
_d = 5
_W = 0.8 * jax.random.normal(_k1, (_d, _d))
_z0 = jax.random.normal(_k2, (_d,))
_zT_ref = _dopri(lambda z, t, W: jnp.tanh(W @ z), _z0, jnp.array([0.0, 1.0]), _W, rtol=1e-7, atol=1e-7)[-1]
assert np.allclose(odeint_unrolled(_z0, _W), _zT_ref, atol=1e-5), "z(T) from odeint_unrolled doesn't match a high-accuracy adaptive integrator"
assert np.allclose(odeint_adjoint(_z0, _W), _zT_ref, atol=1e-5), "Forward z(T) of odeint_adjoint is wrong"
_L_u = lambda z0, W: jnp.sum(odeint_unrolled(z0, W) ** 2)
_L_a = lambda z0, W: jnp.sum(odeint_adjoint(z0, W) ** 2)
_gu = jax.grad(_L_u, argnums=(0, 1))(_z0, _W)
_ga = jax.jit(jax.grad(_L_a, argnums=(0, 1)))(_z0, _W)
assert np.allclose(_ga[0], _gu[0], atol=1e-4), "The adjoint ∂L/∂z0 doesn't match differentiating the unrolled version (check the backward integration of a)"
assert np.allclose(_ga[1], _gu[1], atol=1e-4), "The adjoint ∂L/∂W doesn't match differentiating the unrolled version (check the integral for W̄)"
_n_adj = _resid_elems(_L_a, _z0, _W)
assert _n_adj < 100 * _d, f"The adjoint method shouldn't save the trajectory: the backward pass saved {_n_adj} elements (a 100-step × {_d}-dim trajectory is {100 * _d})"
print(f"Residual elements saved by the adjoint method: {_n_adj} (unrolled version: {_resid_elems(_L_u, _z0, _W)})")
print("✅ Extra exercise C: self-check passed")
Extra exercise D: a vector norm that is safe at 0 (custom_jvp)¶
$\|x\|_2$ is not differentiable at $x=0$, and jax.grad(jnp.linalg.norm) gives nan there. Implement safe_norm with custom_jvp: the forward pass is $\|x\|$, the JVP is $\dfrac{x^\top t}{\|x\|}$, and by convention it is 0 at $x=0$. Requirements:
gradis all zeros at 0 and matchesjnp.linalg.normat nonzero points;- Both forward mode (
jvp) and reverse mode work, and it passescheck_grads.
(Compare with the safe_sqrt written with "double where" in Basics 03: Automatic differentiation with grad: here custom_jvp achieves the same goal with a more direct rule.)
Self-check interface: safe_norm(x): defined with jax.custom_jvp; x is a 1D array, and it returns the scalar $\|x\|_2$.
# ✍️ Your answer: Extra exercise D
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise D
try:
safe_norm
except NameError as e:
print(f"⏳ Extra exercise D: not answered yet: {e}")
else:
assert isinstance(safe_norm, jax.custom_jvp), "safe_norm should be defined with jax.custom_jvp"
_z = jnp.zeros(5)
_g0 = jax.jit(jax.grad(safe_norm))(_z)
assert np.all(np.isfinite(_g0)) and np.all(_g0 == 0.0), f"The gradient at x = 0 should be all zeros; got {_g0}"
_x = jax.random.normal(jax.random.key(11), (5,))
_x64 = np.asarray(_x, np.float64)
assert np.isclose(float(safe_norm(_x)), np.linalg.norm(_x64), rtol=1e-6), "The forward value is not the L2 norm"
assert np.allclose(jax.grad(safe_norm)(_x), _x64 / np.linalg.norm(_x64), atol=1e-6), "At nonzero points the gradient should be x / ‖x‖"
_t = jax.random.normal(jax.random.key(12), (5,))
assert np.isclose(float(jax.jvp(safe_norm, (_x,), (_t,))[1]), _x64 @ np.asarray(_t, np.float64) / np.linalg.norm(_x64), rtol=1e-5), "The JVP should be xᵀt / ‖x‖"
assert np.allclose(jax.vmap(jax.grad(safe_norm))(jnp.stack([_z, _x])), np.stack([np.zeros(5), _x64 / np.linalg.norm(_x64)]), atol=1e-6), "Wrong gradient under vmap"
check_grads(safe_norm, (_x,), order=1, modes=["fwd", "rev"])
print("✅ Extra exercise D: self-check passed")
Quiz¶
The answers to all 8 questions below have been verified by running code in this notebook.
1. To get the full Jacobian of $f:\mathbb{R}^{1000}\to\mathbb{R}^{3}$, should you use jacfwd or jacrev? How many JVPs or VJPs does each need?
Answer and explanation
Use jacrev: it needs only 3 VJPs (one per output), while jacfwd needs 1000 JVPs. The general rule: few outputs, use reverse; few inputs, use forward.
2. What happens with jax.linear_transpose(jnp.sin, x)? How is jax.vjp related to linearize and linear_transpose?
Answer and explanation
It raises NotImplementedError, because sin isn't linear and has no transpose rule. vjp = linear_transpose ∘ linearize: first linearize to get a linear function of the tangent $v\mapsto Jv$, then transpose it to get $u\mapsto J^\top u$ (see Section 4.1).
3. What does this code output? jax.jvp(predict, (params, inputs), (tangents, jax.lax.stop_gradient(inputs)))
Answer and explanation
The value of stop_gradient(inputs) is just inputs itself, so it's treated as a nonzero tangent, and the output gains an extra inputs @ w term (in Section 4.2 the result is [100, 130] instead of [30, 30]). To say "don't perturb", pass jnp.zeros_like(inputs). stop_gradient affects differentiation of its output; it doesn't make its value zero.
4. If a function's gradient is defined with jax.custom_vjp, can you use jax.hessian on it? What about with jax.custom_jvp?
Answer and explanation
Not with custom_vjp: the outer layer of hessian = jacfwd(jacrev(f)) is forward mode, and custom_vjp doesn't support JVPs, so you get TypeError. With custom_jvp it works: forward, reverse (by automatic transposition) and higher-order derivatives are all supported. The conditions are that the rule is linear in the tangent (the tangent side uses only JAX primitives) and that the rule itself is differentiable.
5. What happens when jax.grad is applied to lax.while_loop? Does adding jax.checkpoint to the loop body fix it? What are the right alternatives?
Answer and explanation
It raises ValueError: Reverse-mode differentiation does not work for lax.while_loop .... Adding checkpoint doesn't fix it. Alternatives: when the number of iterations is known, use fori_loop with static bounds or scan; for iterate-until-convergence, use custom_vjp + the implicit function theorem; or use forward mode only (jvp/jacfwd). Also note that inside jit, a fori_loop whose upper bound is a tracer also falls back to while_loop.
6. Is the gradient of jax.grad(lambda x: (jnp.sort(x) * w).sum()) zero? What about argmax?
Answer and explanation
The gradient of sort isn't zero: sort is an input-dependent permutation, and the gradient flows back through that permutation, e.g. giving [1, 3, 2] in Section 4.9. The output of argmax is an integer, so calling grad on it directly raises TypeError; after casting to float the gradient is 0. When you need a differentiable alternative, use soft-argmax or Gumbel-softmax (with STE).
7. What is perturbation confusion? Why doesn't it happen in JAX?
Answer and explanation
In nested forward AD, if the infinitesimals $\varepsilon$ of different levels aren't kept apart, the inner derivative mistakes the outer perturbation for its own and the result is wrong (in the Section 4.4 example the naive implementation gets 2, while the correct answer is 1). Every jvp/grad in JAX creates a new trace with a unique level, every tracer remembers which trace it belongs to, and an inner trace treats outer tracers as constants, so there's no confusion.
8. When differentiating a fixed point $x^\star=T(a,x^\star)$ with custom_vjp, what does the backward pass need to save? What is the adjoint equation? What's the advantage over "unroll the iterations and backprop"?
Answer and explanation
Only $x^\star$ (and $a$) needs to be saved. The adjoint equation is $u = g + (\partial T/\partial x)^\top u$ (itself a fixed-point problem), and the gradient is $\bar a = u^\top \partial T/\partial a$. Three advantages: memory doesn't depend on the number of iterations; the forward pass can iterate to convergence with while_loop (the unrolled version can't do reverse mode through that); and the gradient's accuracy depends only on the accuracy of the solution, not on the iteration path.
Version notes¶
Below is a summary of the patterns and conclusions in this chapter that need attention under JAX 0.12. Each one is verified with an assert at the corresponding place in the text.
| No. | Section | Key point | Where |
|---|---|---|---|
| 1 | 4.2 | The forward output inputs @ w + b in the PyTree example is [70.1, 99.9] (not 100.9) | 4.2 Example 2 ✅ cell |
| 2 | 4.2 | The value of stop_gradient(inputs) is inputs itself, so it can't serve as a zero tangent (it turns the tangent output into [100,130]); passing None raises TypeError; to not perturb, pass jnp.zeros_like(inputs) or use a closure | 4.2 ❌/✅ cells |
| 3 | 4.4 | The HVP written as grad over grad is reverse-over-reverse, not "reverse-over-forward" | 4.4 HVP explanation |
| 4 | 4.4/4.5 | In the source, jax.hessian(f) is always jacfwd(jacrev(f)) | 4.5 Hessian explanation |
| 5 | 4.5 | In jax.jvp(grad_f, (x_in,), (v_)), the undefined v_ raises NameError; the tangent must be the tuple (v,) | 4.5 ❌/✅ cells |
| 6 | 4.6 | stable_log1pexp with log1p(exp(x)) in the forward pass only has a stable gradient; the forward pass is inf at $x=100$; a stable forward pass uses logaddexp(x,0) | 4.6 Example 2 |
| 7 | 4.6 | Unstable gradients raise no exception by default and silently return nan, so try/except FloatingPointError catches nothing; only jax.debug_nans(True) makes it raise | 4.6 Example 2 |
| 8 | 4.6 | Don't create a key if you don't need random numbers; when you do, use jax.random.key rather than the old-style jax.random.PRNGKey | 4.6 Example 1 ❌/✅ cells |
| 9 | 4.6 | Non-differentiable array arguments can return None in bwd (treated as 0); non-array arguments need nondiff_argnums (the bwd signature becomes bwd(*nondiff, res, g), and these arguments can't be tracers) | 4.6 extra |
| 10 | 4.7 | The safe_x trick in custom_jvp belongs only in the derivative formula; using safe_x in the forward pass too makes stable_sqrt(0) return 1.0, so the forward pass should just use jnp.sqrt(x) | 4.7 ❌/✅ cells |
| 11 | 4.7 | Any JVP rule that is linear in the tangent, with the tangent side using only JAX primitives, can be transposed automatically into a VJP; custom_jvp and custom_vjp can't both be defined (custom_jvp has no defvjp); custom_vjp doesn't support forward mode | 4.7 explanation of the relation to custom_vjp |
| 12 | 4.8 | while_loop only supports forward mode; grad raises ValueError, and adding checkpoint doesn't help; use a static fori_loop/scan or the implicit function theorem instead | 4.8 ❌/✅ cells |
| 13 | 4.8 | In the scan RNN example, differentiating with respect to $h_0$ gives the gradient of the initial state, not a parameter gradient; the function is named grad_rnn_h0 | 4.8 scan example |
| 14 | 4.9 | The error type when no differentiation rule exists depends on the case: TypeError (integer inputs/outputs), ValueError (pure_callback, while_loop), NotImplementedError (missing transpose rule) | 4.9 explanation of why differentiation fails |
| 15 | 4.10 | Naive softplus gives inf in the forward pass and nan as the gradient at $x=100$; at $x=88$ the gradient silently becomes 0.0 because a subnormal number is flushed to zero, and XLA doesn't rewrite it into a stable form | 4.10 opening explanation |
| 16 | 4.10 | Naive softplus doesn't raise OverflowError when it overflows (the except branch is dead code); the actual result is nan | 4.10 ❌/✅ cells |
| 17 | 4.10 | For outputs not on the gradient path, the g received by bwd is an all-zero array by default, not None; it's a symbolic zero only with symbolic_zeros=True | 4.10 considerations |
| 18 | 4.10 | The forward pass of softplus_custom is still the naive implementation; a complete implementation uses a stable forward pass (logaddexp) plus custom_vjp, verified with check_grads | 4.10 ✅ cell |
Chapter summary¶
- What the two modes really are: a JVP computes $Jv$ at a cost independent of the output dimension; a VJP computes $u^\top J$ at a cost independent of the input dimension. In JAX, VJP = linearize + transpose (
linearize+linear_transpose): each primitive only provides a JVP rule, and linear primitives also provide a transpose rule. - Higher-order derivatives: transformations nest freely (for HVPs, prefer forward-over-reverse); JAX avoids perturbation confusion with level-tagged traces; Taylor mode in
jax.experimental.jetstill works in 0.12 and gives several derivative orders in one pass, but doesn't support control flow orcustom_jvp. - Full matrices:
jacfwd/jacrevare essentiallyvmapwrapped around JVP/VJP;hessian = jacfwd(jacrev); for sparse Jacobians use coloring, which cuts the number of JVPs from $n$ to the number of colors. custom_vjp:fwdreturns(out, res), andbwd(res, g)returns a tuple as long as the argument list; non-array arguments go innondiff_argnums(and can't be tracers). Typical uses are gradient clipping, gradient reversal, STE, and the implicit function theorem (fixed points, DEQ, ODE adjoints). No forward mode support.custom_jvp: usedefjvpordefjvps; as long as the rule is linear in the tangent (and the tangent side uses only JAX primitives), it gets reverse mode automatically; it can't be combined withdefvjpon the same function. Prefer it whenever you can.- Control flow:
cond,scanand staticfori_loopcan be differentiated in both directions;while_looponly supports forward mode;checkpoint+scanreduces backward-pass memory from "every step's intermediates" to "every step's carry" (quantitative comparison in Advanced Chapter 6 §6.6). - Non-differentiable ops:
round/floor/sign/argmaxhave zero gradients (callinggradon an integer output raises);sort/top_k/maxpass gradients back through the permutation or selection;stop_gradientis "identity forward, zero derivative"; smooth approximations (sigmoid, soft-argmax, Gumbel-softmax) or STE can stand in. - Verification: cross-check every custom rule with
check_gradsor an analytic solution; numerical problems (inf/nan/subnormals flushed to zero) don't raise exceptions, so turn onjax.debug_nanswhen needed.
Cheat sheet¶
| API | Purpose |
|---|---|
jax.jvp(f, primals, tangents) | Forward mode: returns (f(x), J v); tangents must match the primals in dtype and structure |
jax.vjp(f, *primals) | Reverse mode: returns (f(x), vjp_fun), and vjp_fun(u) returns $u^\top J$ (a tuple) |
jax.linearize(f, x) | Returns (f(x), f_lin) with f_lin(v) = J v; residuals are computed once |
jax.linear_transpose(f_lin, x) | Transpose of a linear function; vjp = linear_transpose ∘ linearize |
jax.jacfwd / jax.jacrev | Full Jacobian: vmap(jvp) / vmap(vjp) |
jax.hessian(f) | jacfwd(jacrev(f)) |
jax.jvp(jax.grad(f), (x,), (v,))[1] | HVP (forward-over-reverse) |
jax.experimental.jet.jet(f, (x,), ((1., 0., ...),)) | Taylor mode: derivatives of order 1~k in one go (no control flow or custom_jvp) |
vmap(lambda s: jvp(f,(x,),(s,))[1])(seeds) | Sparse Jacobian via coloring |
@jax.custom_vjp + f.defvjp(fwd, bwd) | Custom backward: fwd → (out, res), bwd(res, g) → tuple |
partial(jax.custom_vjp, nondiff_argnums=(i,)) | Non-array arguments; bwd(*nondiff, res, g) |
defvjp(..., symbolic_zeros=True) | bwd receives symbolic zeros and can skip computation |
@jax.custom_jvp + @f.defjvp | Custom forward rule(primals, tangents) → (out, t_out), reverse mode comes for free |
f.defjvps(jvp_x, None, ...) | Per-argument JVP rules jvp_i(t_i, ans, *primals) |
jax.test_util.check_grads(f, args, order, modes) | Cross-check against finite differences; custom_vjp can only use modes=['rev'] |
lax.cond / lax.scan / lax.fori_loop (static bounds) | Both forward and reverse supported |
lax.while_loop | Forward mode only |
jax.checkpoint(step) inside scan | Recompute in-step intermediates in the backward pass to save memory |
jax.ad_checkpoint.print_saved_residuals(f, *args) | See which tensors the backward pass saves |
jit(f).lower(*a).compile().memory_analysis() | See compiled temporary memory (usage in Advanced Chapter 6 §6.6) |
lax.stop_gradient(x) | Identity forward, zero derivative |
x + stop_gradient(round(x) - x) | One-line STE |
jax.random.gumbel + softmax((logits+g)/τ) | Gumbel-softmax |
jax.debug_nans(True) | Raise FloatingPointError as soon as a NaN appears |
Appendix: reference solutions¶
First write your own answer in the "✍️ Your answer" cell in the main text and check it with the "🧪 Self-check"; come here to compare once the self-check passes or if you're really stuck. Running the solutions below in order overwrites any functions with the same names you defined in the answer cells.
Course exercise 1¶
See "Course exercise 1: compare the naive gradient with the custom gradient" in the main text.
# Reference solution: Course exercise 1
xs = jnp.linspace(-20.0, 20.0, 41)
g_naive = jax.vmap(jax.grad(softplus_naive))(xs)
g_custom = jax.vmap(jax.grad(softplus_custom))(xs)
print("Max difference:", float(jnp.max(jnp.abs(g_naive - g_custom))))
assert jnp.allclose(g_naive, g_custom, atol=1e-6)
Course exercise 2¶
See "Course exercise 2: choosing residuals — save or recompute" in the main text.
# Reference solution: Course exercise 2
@jax.custom_vjp
def sin2_save(x):
return jnp.sum(jnp.sin(x) ** 2)
sin2_save.defvjp(lambda x: (jnp.sum(jnp.sin(x) ** 2), (jnp.sin(x), jnp.cos(x))),
lambda r, g: (g * 2 * r[0] * r[1],))
@jax.custom_vjp
def sin2_recompute(x):
return jnp.sum(jnp.sin(x) ** 2)
sin2_recompute.defvjp(lambda x: (jnp.sum(jnp.sin(x) ** 2), x),
lambda x, g: (g * jnp.sin(2 * x),)) # 2 sin cos = sin 2x, recomputed in bwd
xv = jnp.linspace(0.0, 3.0, 1000)
print("A (saves sin, cos):"); print_saved_residuals(sin2_save, xv)
print("B (saves only x) :"); print_saved_residuals(sin2_recompute, xv)
assert jnp.allclose(jax.grad(sin2_save)(xv), jax.grad(sin2_recompute)(xv), atol=1e-6)
assert jnp.allclose(jax.grad(sin2_recompute)(xv), jnp.sin(2 * xv), atol=1e-6)
Course exercise 3¶
See "Course exercise 3: softplus for production" in the main text.
# Reference solution: Course exercise 3
softplus_prod = softplus_stable # the stable implementation from Section 4.10 already meets the requirements; verify it on an extreme range below
xs = jnp.array([-1000.0, -50.0, 0.0, 50.0, 1000.0])
y = jax.jit(jax.vmap(softplus_stable))(xs)
gy = jax.jit(jax.vmap(jax.grad(softplus_stable)))(xs)
print(y, gy)
assert jnp.all(jnp.isfinite(y)) and jnp.all(jnp.isfinite(gy))
assert jnp.allclose(y, np.logaddexp(np.asarray(xs, np.float64), 0.0))
Extra exercise A¶
See "Extra exercise A: global-norm gradient clipping for a pytree" in the main text.
# Reference solution: Extra exercise A
@partial(jax.custom_vjp, nondiff_argnums=(1,))
def clip_grad_global_norm(params, max_norm):
return params
def _cgn_fwd(params, max_norm):
return params, None
def _cgn_bwd(max_norm, _, g):
norm = jnp.sqrt(sum(jnp.sum(l ** 2) for l in jax.tree.leaves(g)))
scale = jnp.minimum(1.0, max_norm / (norm + 1e-12))
return (jax.tree.map(lambda l: l * scale, g),)
clip_grad_global_norm.defvjp(_cgn_fwd, _cgn_bwd)
def loss_p(p, x):
return jnp.sum((x @ p["w"] + p["b"]) ** 2)
p = {"w": jnp.ones((3, 2)), "b": jnp.array([0.5, -0.5])}
xb = jnp.arange(6.0).reshape(2, 3)
g_raw = jax.grad(loss_p)(p, xb)
g_clip = jax.grad(lambda p, x: loss_p(clip_grad_global_norm(p, 1.0), x))(p, xb)
gnorm = lambda t: float(jnp.sqrt(sum(jnp.sum(l ** 2) for l in jax.tree.leaves(t))))
print("Original norm:", gnorm(g_raw), " after clipping:", gnorm(g_clip))
assert abs(gnorm(g_clip) - 1.0) < 1e-5
assert all(jnp.allclose(a / gnorm(g_raw), b) for a, b in zip(jax.tree.leaves(g_raw), jax.tree.leaves(g_clip)))
# Cross-check against optax.clip_by_global_norm
import optax
tx = optax.clip_by_global_norm(1.0)
g_optax, _ = tx.update(g_raw, tx.init(p))
assert all(jnp.allclose(a, b, atol=1e-6) for a, b in zip(jax.tree.leaves(g_optax), jax.tree.leaves(g_clip)))
Extra exercise B¶
See "Extra exercise B: coloring a pentadiagonal Jacobian" in the main text.
# Reference solution: Extra exercise B
def f_penta(x):
pad = jnp.pad(x, 2)
xm2, xm1, xp2 = pad[:-4], pad[1:-3], pad[4:]
return x**3 + xm2 * xp2 - xm1
def colored_jacobian_banded(f, x, bandwidth):
n, c = x.shape[0], 2 * bandwidth + 1
j = jnp.arange(n)
seeds = (j[None, :] % c == jnp.arange(c)[:, None]).astype(x.dtype)
comp = jax.vmap(lambda s: jax.jvp(f, (x,), (s,))[1])(seeds) # only c JVPs
ii, jj = j[:, None], j[None, :]
return jnp.where(jnp.abs(ii - jj) <= bandwidth, comp[jj % c, ii], 0.0)
x50 = jnp.linspace(-1.0, 1.0, 50)
J5 = colored_jacobian_banded(f_penta, x50, 2)
assert jnp.allclose(J5, jax.jacfwd(f_penta)(x50))
# Coloring the same function as if tridiagonal (not enough colors) goes wrong: the number of colors must be >= 2b+1
J3_wrong = colored_jacobian_banded(f_penta, x50, 1)
assert not jnp.allclose(J3_wrong, jax.jacfwd(f_penta)(x50))
print("5-color coloring is correct; with only 3 colors the decompressed result is wrong (columns collide)")
Extra exercise C¶
See "Extra exercise C: differentiating an ODE with the adjoint method" in the main text.
# Reference solution: Extra exercise C
def ode_f(z, W):
return jnp.tanh(W @ z)
def rk4_step(f, s, h, W):
add = lambda s, k, c: jax.tree.map(lambda a, b: a + c * b, s, k)
k1 = f(s, W)
k2 = f(add(s, k1, h / 2), W)
k3 = f(add(s, k2, h / 2), W)
k4 = f(add(s, k3, h), W)
return jax.tree.map(lambda s_, a, b, c, d: s_ + h / 6 * (a + 2 * b + 2 * c + d), s, k1, k2, k3, k4)
N_STEPS, T_END = 100, 1.0
def odeint_unrolled(z0, W):
h = T_END / N_STEPS
return lax.fori_loop(0, N_STEPS, lambda i, z: rk4_step(ode_f, z, h, W), z0)
@jax.custom_vjp
def odeint_adjoint(z0, W):
return odeint_unrolled(z0, W)
def _ode_fwd(z0, W):
zT = odeint_unrolled(z0, W)
return zT, (zT, W) # save only the endpoint, not the trajectory
def _ode_bwd(res, g):
zT, W = res
def aug_dynamics(s, W):
z, a, _ = s
_, vjp = jax.vjp(ode_f, z, W)
a_z, a_W = vjp(a)
return (ode_f(z, W), -a_z, -a_W) # d/dt (z, a, W̄)
h = -T_END / N_STEPS # integrate backward
s = (zT, g, jnp.zeros_like(W))
_, a0, gW = lax.fori_loop(0, N_STEPS, lambda i, s: rk4_step(aug_dynamics, s, h, W), s)
return a0, gW
odeint_adjoint.defvjp(_ode_fwd, _ode_bwd)
k1, k2 = jax.random.split(jax.random.key(1))
W_ode = jax.random.normal(k1, (4, 4))
z0_ode = jax.random.normal(k2, (4,))
L_unr = lambda z0, W: jnp.sum(odeint_unrolled(z0, W) ** 2)
L_adj = lambda z0, W: jnp.sum(odeint_adjoint(z0, W) ** 2)
gz_u, gW_u = jax.grad(L_unr, argnums=(0, 1))(z0_ode, W_ode)
gz_a, gW_a = jax.jit(jax.grad(L_adj, argnums=(0, 1)))(z0_ode, W_ode)
print("dL/dz0 max difference:", float(jnp.max(jnp.abs(gz_u - gz_a))), " dL/dW max difference:", float(jnp.max(jnp.abs(gW_u - gW_a))))
assert jnp.allclose(gz_u, gz_a, atol=1e-4)
assert jnp.allclose(gW_u, gW_a, atol=1e-4)
Extra exercise D¶
See "Extra exercise D: a vector norm that is safe at 0 (custom_jvp)" in the main text.
# Reference solution: Extra exercise D
@jax.custom_jvp
def safe_norm(x):
return jnp.sqrt(jnp.sum(x ** 2))
@safe_norm.defjvp
def _safe_norm_jvp(primals, tangents):
x, = primals
t, = tangents
n = safe_norm(x)
safe_n = jnp.where(n == 0.0, 1.0, n) # avoid 0/0
return n, jnp.where(n == 0.0, 0.0, jnp.dot(x, t) / safe_n)
z3 = jnp.zeros(3)
print("Gradient of jnp.linalg.norm at 0:", jax.grad(jnp.linalg.norm)(z3))
print("Gradient of safe_norm at 0 :", jax.grad(safe_norm)(z3))
assert jnp.all(jnp.isnan(jax.grad(jnp.linalg.norm)(z3)))
assert jnp.all(jax.grad(safe_norm)(z3) == 0.0)
xv = jnp.array([3.0, -4.0, 12.0])
assert jnp.allclose(jax.grad(safe_norm)(xv), jax.grad(jnp.linalg.norm)(xv))
assert jnp.allclose(jax.jvp(safe_norm, (xv,), (jnp.ones(3),))[1], jnp.sum(xv) / 13.0)
check_grads(safe_norm, (xv,), order=1, modes=["fwd", "rev"])
# Second order: in float32, second-order finite differences are noisy (measured relative error about 1%), so cross-check against the analytic Hessian instead
# ∇²||x|| = (I - x x^T / ||x||^2) / ||x||
H_ref = (np.eye(3) - np.outer(xv, xv) / 169.0) / 13.0
assert np.allclose(jax.hessian(safe_norm)(xv), H_ref, atol=1e-6)
print("safe_norm passes check_grads (fwd + rev), and its Hessian matches the analytic solution")