Chapter 3: Automatic Differentiation with grad

Personal study notes on Chapter 3 of the course "Getting Started with JAX" (ApX Machine Learning) · Environment: JAX 0.12 (CPU) Original course: https://apxml.com/courses/getting-started-with-jax/chapter-3-automatic-differentiation-grad

Training a machine learning model almost always means computing the gradient of a loss function with respect to the parameters. Automatic differentiation (AD) is the most efficient and most accurate way to compute gradients. This chapter introduces JAX's core transformation jax.grad: give it an ordinary Python function on numeric inputs, and you get back a function that computes its gradient.

Learning goals

1. Compute gradients of scalar functions with jax.grad;
2. Understand the reverse-mode automatic differentiation behind grad;
3. Control which arguments to differentiate with argnums;
4. Compose grad to get higher-order derivatives;
5. Get the function value and the gradient together with jax.value_and_grad;
6. Understand how control flow interacts with differentiation, and the limits of grad.

Sections

How to read this: the code examples follow the course pages. All of them were run under JAX 0.12, and the output below each cell is the actual result.

• The text uses the recommended JAX 0.12 style throughout. Common pitfalls are marked with > ⚠️ **Note**. When code is involved, the outdated or wrong version goes in a ❌ cell, followed right away by the ✅ recommended version, and the conclusion is checked with assert.
• Cells marked "🔍 Extra" are additional material.
• These notes are collected in the Version notes at the end of the chapter.

The explanations are rewritten around the course's key points. The code examples keep the course's variable names and values where possible, and all of them have been run.

3.1 Understanding gradients¶

What is a gradient¶

Start with a function of one variable, $f(x)$. Its derivative at a point is the slope of the tangent line there. It tells you how fast $f$ changes when $x$ changes a little.

For example, $f(x)=x^2$ has $f'(x)=2x$. At $x=2$ the tangent slope is $4$, which is the derivative (gradient) at that point.

For a function with several inputs and one scalar output, $f(x_1,\dots,x_n)$, stack the partial derivatives with respect to each input into a vector. That vector is the gradient:

$$ \nabla f(x_1,\dots,x_n)=\left[\frac{\partial f}{\partial x_1},\ \frac{\partial f}{\partial x_2},\ \dots,\ \frac{\partial f}{\partial x_n}\right] $$

The gradient is a vector. It points in the direction in which the function increases fastest at that point, and its length tells you how steep that increase is.

Example: $f(x,y)=x^2+\sin(y)$

$$ \frac{\partial f}{\partial x}=2x,\qquad \frac{\partial f}{\partial y}=\cos(y),\qquad \nabla f(x,y)=[2x,\ \cos y],\qquad \nabla f(1,0)=[2,\ 1] $$

Why gradients matter for optimization¶

Since the gradient points in the direction of fastest increase, the negative gradient points in the direction of fastest decrease. Training a model means making the loss as small as possible, and gradient descent uses exactly this:

1. Start from some initial parameters;
2. Compute the gradient of the loss at the current parameters;
3. Take a small step along the negative gradient;
4. Repeat steps 2–3 until convergence.

The update rule:

$$ \theta_{\text{new}} = \theta_{\text{old}} - \eta\,\nabla f(\theta_{\text{old}}) $$

Here $\eta$ is the learning rate, a small positive number that controls the step size.

Why we need automatic differentiation¶

Deep networks easily have millions of parameters. Deriving gradients by hand is impractical and error-prone. There are three main ways to let a computer take derivatives:

JAX's grad uses reverse-mode automatic differentiation, which machine learning calls backpropagation. It is a great fit for functions with "many inputs and one scalar output (the loss)".

🔍 Extra: what exactly is wrong with numerical differentiation¶

⚪ Optional: read this if you wonder why we don't train with finite differences, or if you need to write your own gradient check.

The error of numerical differentiation has two sources. If $h$ is too large, you get truncation error ($O(h)$). If $h$ is too small, floating-point subtraction brings in rounding error (about $\varepsilon/h$). JAX uses float32 by default (machine epsilon $\varepsilon\approx 1.2\times10^{-7}$), which makes this worse. Below we differentiate $\sin$ at $x=1$ and compare forward differences, central differences and jax.grad:

In float32, the finite-difference error first drops as $h$ shrinks, then climbs back up because of rounding error. At best it reaches about $10^{-4}\sim10^{-5}$. The autodiff error is at the level of machine precision. So numerical differentiation is only good for checking gradients (see check_grads in 3.8), not for training.

3.2 Introducing jax.grad¶

jax.grad is itself a function transformation. You give it a Python numeric function that returns a scalar, and it gives you back a new function that computes the gradient of the original one.

Take $f(x)=x^3$. Its derivative is $f'(x)=3x^2$, so $f'(2.0)=3\times 2.0^2=12.0$. First define and test the original function:

Use jax.grad to get the gradient function. Note that grad_f, like f, is just an ordinary Python callable:

Key points

• Input: a Python function that returns a single scalar;
• jax.grad returns a new function;
• The new function takes the same arguments as the original and returns the gradient of the output with respect to the first positional argument.

Functions of several arguments: $g(x,y)=x^2 y$. By default jax.grad(g) differentiates only with respect to the first argument $x$: $\dfrac{\partial g}{\partial x}=2xy$, which is $12$ at $(2,3)$. Section 3.4 shows how to differentiate with respect to other arguments.

Requirements on the function

• The function being differentiated must return a single scalar. Returning something non-scalar, like an array or a tuple, raises a TypeError. For functions with several outputs, use jax.jacrev / jax.jacfwd (jax.jacobian is an alias of jax.jacrev), or the lower-level jax.vjp / jax.jvp. Like grad, they are all built on jax.vjp (reverse) or jax.jvp (forward); see 3.3.
• The inputs you differentiate with respect to must be floating-point (or complex). Passing an integer raises a TypeError right away; JAX will not convert it to float for you. If you really need an integer argument (an index, a step count, and so on), put it in a position you don't differentiate, that is, leave it out of argnums. If you must list it in argnums, add allow_int=True; its gradient is then a zero of type float0.

⚠️ Note: it's easy to assume that multi-output tools like jax.jacobian are built on top of jax.grad. They aren't. The source of jax.jacobian simply calls jacrev (the docs say Alias of jax.jacrev), and both jacrev and grad are implemented on top of jax.vjp.

⚠️ Note: using an integer input is not just "not recommended". It raises TypeError: grad requires real- or complex-valued inputs.

One more example: $h(x)=\sin x$, $h'(x)=\cos x$. At $\pi/2$ the mathematical value of the gradient is $\cos(\pi/2)=0$.

⚠️ Note: grad_h(jnp.pi / 2.0) does not print 0.0. It prints about -4.37e-08 (float32 by default). The reason is that $\pi/2$ gets rounded when converted to float32, and cos at the rounded point is not exactly 0. Compare floating-point results with jnp.allclose, not ==.

3.3 How autodiff works: reverse mode¶

Automatic differentiation (AD) is the technique behind jax.grad. It breaks a program into elementary operations and applies the chain rule to them one by one, which gives exact derivatives. AD comes in two flavors: forward mode and reverse mode. jax.grad uses reverse mode, usually called backpropagation in machine learning.

The computation graph¶

Think of a computation as a chain of elementary operations. You can draw it as a computation graph of nodes (operations) and directed edges (data flow). Take $f(x)=\sin(x^2)$ and split it into two steps:

$$ a = x^2,\qquad y=\sin(a) $$

  x ──[square]──► a ──[sin]──► y

Forward pass: computing values¶

Take $x=2.0$:

$$ a = 2.0^2 = 4.0,\qquad y=\sin(4.0)\approx -0.757 $$

During the forward pass, the AD system stores the intermediate values (here $x$ and $a$) and the structure of the graph for the backward pass.

Backward pass: computing gradients with the chain rule¶

Write $\bar v = \dfrac{\partial y}{\partial v}$ (the "sensitivity" of the output $y$ to a variable $v$, also called the adjoint). Walk back from the output:

1. Initialize: the sensitivity of the output to itself is $\bar y=\dfrac{dy}{dy}=1$.
2. Through the sin node: $\bar a=\bar y\cdot\dfrac{dy}{da}=\bar y\cos(a)=\cos(4.0)\approx-0.654$.
3. Through the square node: $\bar x=\bar a\cdot\dfrac{da}{dx}=\bar a\cdot 2x\approx-0.654\times4.0\approx-2.615$.

$\bar x$ is the $\dfrac{dy}{dx}$ we want. Put simply, backpropagation starts from the output and applies the chain rule step by step, getting the sensitivity of the output to each intermediate variable and finally to each input.

Forward (values):     x=2.0 ──► a=4.0 ──► y≈-0.757      (store x, a)
Backward (gradients): x̄≈-2.615 ◄── ā≈-0.654 ◄── ȳ=1

Why machine learning uses reverse mode¶

A loss function $L=f(W_1,\dots,W_n,x,y)$ usually has millions of inputs (parameters) but only one scalar output.

• Reverse mode: one backward pass gives the gradient with respect to all inputs. The cost is a small constant multiple of evaluating the original function (usually quoted as 2–4x), and it barely depends on the number of inputs.
• Forward mode: each pass gives the derivative along only one input direction. Getting the full gradient takes a number of passes proportional to the number of inputs, which is far too expensive for large models.

What JAX does behind the scenes¶

1. Tracing: JAX runs the function once with "tracer" objects in place of the inputs and records the sequence of operations (the computation graph). With grad alone, the tracers carry concrete values, so the forward values are computed on the spot while the operations needed for differentiation are recorded. Only inside jit are the tracers abstract values that carry just shape and dtype, and what gets recorded is the intermediate representation jaxpr (3.7 explains the difference).
2. Forward pass: run these operations and compute the intermediate values the gradient needs (with value_and_grad you also get the function value along the way).
3. Backward pass: from the recorded graph, automatically generate and run the backward computation to get the gradient with respect to the chosen inputs.

You only write the forward computation in NumPy-like syntax. Deriving and computing the gradient is left to JAX.

⚠️ Note: with grad alone, tracers are not abstract. A tracer prints as GradTracer(primal=2.0, ...) and carries a concrete value. Abstract tracers that carry only shape and dtype are a jit behavior (a cell in 3.7 checks this with assert).

🔍 Extra: look at the jaxpr yourself, and how jvp/vjp relate to the two modes¶

🟢 Must-read

(1) See what grad generates. jax.make_jaxpr prints the jaxpr produced by tracing. Compare the jaxpr of the original function with that of the gradient function, and you'll see the backward computation contains cos (the derivative of sin) and a multiplication by 2x:

(2) Forward mode = jax.jvp, reverse mode = jax.vjp. Each of the two modes above has a low-level API in JAX. Let $f:\mathbb R^n\to\mathbb R^m$ with Jacobian $J\in\mathbb R^{m\times n}$:

• jax.grad(f)(x) is essentially vjp_fn(1.0) (when $m=1$, $J$ has a single row, so one VJP gives the whole gradient).
• jax.jvp returns the function value and the directional derivative together. Forward mode doesn't need to store intermediate values, so it uses little memory.
• jax.vjp returns "the function value plus a vjp_fn you can call many times". The intermediate values stored during the forward pass live inside vjp_fn.

🔍 Extra: choosing between jacfwd and jacrev (decision rule + timing check)¶

⚪ Optional: read this if you need a full Jacobian or Hessian and care about efficiency.

To compute the full Jacobian $J\in\mathbb R^{m\times n}$ of $f:\mathbb R^n\to\mathbb R^m$:

• jacfwd uses forward mode and builds it column by column, with $n$ JVPs (JAX batches them in one go with vmap);
• jacrev uses reverse mode and builds it row by row, with $m$ VJPs, and it has to store the forward intermediate values, so it uses more memory.

Both give the same result; they differ only in efficiency. That leads to the following rules:

Let's check this with timings on CPU (absolute times depend on the machine; only the relative sizes matter).

Note: here the weight matrix is passed as an argument to the jit-compiled function instead of being captured by a closure. That is the style JAX recommends (arrays captured by a closure get embedded in the compiled result as constants). It also sidesteps a known issue with the CPU backend in this environment: "closure-captured large array + matmul + reduction" silently returns wrong results. See 00 · Known issues in this environment (root cause in Advanced 02 §2.3).

The timings show that the two combinations with reverse mode on the inside (jacfwd(jacrev), jacrev(jacrev)) are clearly faster than the ones with forward mode inside, because the inner step computes the gradient of a scalar function ($n\to1$). Between those two the difference is small (which one wins depends on the machine). jax.hessian picks jacfwd(jacrev) because forward mode on the outside doesn't need to store the intermediates that reverse mode needs, so it saves memory.

3.4 Differentiating with respect to arguments¶

By default jax.grad differentiates with respect to the first argument. In machine learning, though, a loss function usually takes both the model parameters and the data, and we only want gradients with respect to the parameters. The argnums argument lets you choose which positional argument (or arguments) to differentiate.

Choosing a single argument¶

Take $f(x,y)=x^2 y$: $\dfrac{\partial f}{\partial x}=2xy$, $\dfrac{\partial f}{\partial y}=x^2$. Use $x=3,\ y=4$:

argnums counts from 0, so argnums=1 means differentiate with respect to the second argument $y$:

Choosing several arguments¶

Pass a tuple of integers to argnums, and the gradient function returns a tuple in the same order as argnums. Arguments that aren't selected are treated as constants during differentiation.

Take the affine function $g(w,b,x)=wx+b$: $\dfrac{\partial g}{\partial w}=x$, $\dfrac{\partial g}{\partial b}=1$.

Connection to machine learning¶

During training we want the gradient of the loss with respect to the parameters, $\nabla_{\text{params}}\,\text{loss}(\text{params},\ \text{data\_batch})$; the data batch is just a fixed input. So the usual convention is to put the parameters first and use the default argnums=0.

Real model parameters are usually nested dicts / lists / tuples, which JAX calls a pytree. jax.grad supports pytrees directly, and the gradient has exactly the same structure as the parameters. So to update the parameters, use jax.tree.map to apply p - lr * g to each leaf (array).

⚠️ Note: updated_params = current_params - learning_rate * gradients only works when the parameters are a single array. When they are a dict it raises TypeError: unsupported operand type(s); use jax.tree.map(lambda p, g: p - learning_rate * g, params, grads) instead.

Here is that example filled out into runnable code:

🔍 Extra: has_aux=True for functions that return more than the loss¶

🟢 Must-read

grad requires the function to return a scalar. But a loss function often wants to hand back auxiliary information too, like predictions or accuracy, and returning a tuple directly raises an error. The right way is to return (loss, aux) and add has_aux=True. The gradient function then returns (grads, aux), and aux is not differentiated.

3.5 Higher-order derivatives (grad of grad)¶

jax.grad takes a function $f$ with scalar output and returns a new Python function that computes the gradient. JAX transformations take Python functions in and give Python functions back, so the result of a transformation can be transformed again. For a scalar function of one variable, grad gives you another scalar function, so you can apply grad once more and get the second derivative.

Take $f(x)=x^3$: $f'(x)=3x^2$, $f''(x)=6x$. At $x=2$: $f=8,\ f'=12,\ f''=12$.

You can also nest the calls directly. The same trick extends to third and higher orders, at the cost of more computation and possibly worse numerical stability:

Functions of several variables and the Hessian¶

For a function of several variables, all the second partial derivatives together form the Hessian matrix:

$$ H_{ij}=\frac{\partial^2 f}{\partial x_i\,\partial x_j} $$

Example: $f(x,y)=x^2y+y^3$.

$$ \frac{\partial f}{\partial x}=2xy,\quad \frac{\partial^2 f}{\partial x^2}=2y,\quad \frac{\partial^2 f}{\partial y\,\partial x}=2x $$

Differentiating with respect to $x$ twice (argnums=0) gives $\partial^2 f/\partial x^2$. Differentiating first with respect to $x$ and then switching to argnums=1 for $y$ gives the mixed partial. At $(2,3)$ they are $6$ and $4$:

For the full Hessian you can use jax.hessian directly, which is usually more efficient.

🔍 Extra: how to use jax.hessian and how it works¶

⚪ Optional: read this if you need second derivatives, Hessians or Hessian-vector products.

• jax.hessian(f) is implemented as jax.jacfwd(jax.jacrev(f)), that is, forward-over-reverse. The inner jacrev computes the gradient (many inputs, one output, a good fit for reverse mode), and the outer jacfwd takes the Jacobian of the gradient ($n$ inputs, $n$ outputs, where forward mode uses less memory).
• jax.hessian differentiates with respect to the first argument by default, so for a function of several variables you usually pack them into one vector. You can also use argnums=(0, 1), which returns a nested tuple ((H_xx, H_xy), (H_yx, H_yy)).
• A full Hessian has $n^2$ entries, which simply won't fit in memory when $n$ is large. Many algorithms (conjugate gradient, Newton-CG) only need Hessian-vector products $Hv$, which you can compute with jax.jvp(jax.grad(f), (x,), (v,)) without ever building the whole matrix (see Exercise C).

3.6 Value and gradient (jax.value_and_grad)¶

When training a model you usually want both the gradient (to update the parameters) and the loss value (to monitor progress). If you call f(x) and jax.grad(f)(x) separately, the forward computation runs twice. But reverse mode already does a forward pass first, so the function value is computed in that step anyway.

jax.value_and_grad(f) returns a new function that, when called, returns a tuple (value, gradient) and runs the forward computation only once.

Example: $p(x)=x^3+2x^2-3x+1$, $p'(x)=3x^2+4x-3$. At $x=2$: $p(2)=8+8-6+1=11$, $p'(2)=12+8-3=17$.

value_and_grad supports argnums too:

• When argnums is an integer, the returned gradient corresponds to that one argument;
• When argnums is a tuple of integers, the returned gradient is also a tuple, with one element per chosen argument.

Example: $m(a,b)=a^2\sin b$, $\dfrac{\partial m}{\partial a}=2a\sin b$, $\dfrac{\partial m}{\partial b}=a^2\cos b$. With $a=3,\ b=\pi/2$: the value is $9$, the gradient with respect to $a$ is $6$, and the gradient with respect to $b$ is mathematically $0$.

⚠️ Note: the gradient with respect to $b$ does not print as 0. It prints about -3.93e-07 (float32). Same reason as in 3.2: $\pi/2$ is rounded in float32, so check it with jnp.allclose.

Optimization algorithms like gradient descent need "loss + gradient" at every step, so value_and_grad is the standard way to write a training loop. It also composes with jax.jit (see the linear regression example in 3.9).

🔍 Extra: the return structure when combined with has_aux¶

🟢 Must-read

In real code the most common combination is jax.value_and_grad(loss_fn, has_aux=True), which returns ((loss, aux), grads). The value part is the (loss, aux) tuple exactly as the function returned it. The parentheses are easy to get wrong when unpacking:

3.7 Differentiation and control flow¶

What grad differentiates is the path the function actually takes for the current input. To understand this, the key is to tell apart two kinds of tracing:

In real training, grad almost always sits inside jit, so value-dependent control flow has to use the structured control-flow primitives jax.lax.cond, jnp.where, jax.lax.fori_loop and jax.lax.scan.

⚠️ Note: it's easy to believe that "JAX traces only once per input shape and type", so grad can only differentiate the one branch taken during tracing, and that a Python if will raise "the truth value of an abstract tracer is ambiguous" at x=-2.0. In fact, both "trace only once" and that error are jit behaviors. With grad alone, a Python if works fine (x=-2.0 returns -1.0); only jit(grad(f))(-2.0) raises TracerBoolConversionError.

Conditionals¶

The course's example: return $x^2$ when $x>0$, otherwise return $-x$.

Rewriting with jax.lax.cond (works under jit too)¶

jax.lax.cond(pred, true_fun, false_fun, *operands) uses pred to decide whether to run true_fun(*operands) or false_fun(*operands). Both branches are traced and the choice is made at run time, so it works with both grad and jit. The two branches must return results with the same shape and dtype. Operands can be passed as several positional arguments (*operands).

For an elementwise choice between two values you can also use jnp.where(pred, a, b), but watch out for its NaN-gradient trap (see the special topic in 3.8).

Loops¶

Static loops: if the number of iterations of a for loop is known at trace time (say, a Python constant), the loop is unrolled into a sequence of ordinary operations, and differentiation works fine. For example, looping 3 times and multiplying by 2 each time gives $y=8x$, whose gradient is $8$. The downside is that with many iterations the unrolled graph gets huge and compilation slows down.

Dynamic loops: if the iteration count or the while condition depends on a traced value, a plain Python loop won't run under jit, and you need the structured control-flow primitives instead. But these primitives differ in how well they support reverse-mode differentiation (grad/vjp):

Reverse mode has to store the intermediate values of every step, and the number of steps of a while_loop isn't known at compile time, so there is no way to allocate that space ahead of time. That's why it isn't supported. For loops you need to differentiate in reverse mode, use lax.scan (fixed count), or declare the count as a static argument (jax.jit(..., static_argnums=...)).

⚠️ Note: neither jax.lax.while_loop nor a fori_loop with traced bounds can be used for loops that need reverse-mode differentiation. Running grad through them raises ValueError: Reverse-mode differentiation does not work for lax.while_loop ...; they only support forward mode. The fori_loop example below works only because n=3 is a concrete integer when grad is used on its own.

body_fun(i, val) returns the new val:

Two final reminders. First, if you don't use these primitives inside jit, you'll get an error at trace time. Second, even without jit, grad only differentiates the path taken for the current input, so gradients near a branch boundary can be misleading.

3.8 Limitations and caveats¶

1. Non-differentiable operations¶

Automatic differentiation relies on the chain rule, which needs every elementary operation to be differentiable. The following operations are either not differentiable in the mathematical sense, or have a derivative that is zero almost everywhere:

• Discontinuous functions: jnp.sign, jnp.round, jnp.floor, jnp.ceil and so on. JAX defines their derivative as 0;
• Casting to integer: x.astype(jnp.int32) is a step function, so the gradient gets cut to 0 there (if the function's output itself is an integer, grad raises a TypeError right away);
• Comparisons: operations like x > 0 produce booleans, which are constant near the input, so the gradient is zero.

At kinks (points where the derivative doesn't exist), JAX assigns each primitive a fixed value that doesn't vary by platform:

Real NaNs mostly come from points like the last two rows, where "the derivative is infinite or 0/0" (inf times 0 also becomes NaN). The special topic below covers this in detail. If you want useful gradients at discontinuities, you have a few options: replace the function with a smooth one (for example, approximate a step with jax.nn.sigmoid(k*x)); use a straight-through estimator (see stop_gradient below); or use an optimization method other than gradient descent.

⚠️ Note: jax.grad(jnp.abs)(0.0) is 1.0, not 0.0.

⚠️ Note: jax.grad(jnp.sign)(0.0) is a fixed 0.0 (the derivative of sign is defined as 0 everywhere). It is never NaN and doesn't vary by platform.

🔍 Extra (special topic): NaN gradients and the "double where" trick¶

🟢 Must-read

The symptom¶

To avoid dangerous points like sqrt(0) or log(0), many people write jnp.where(x > 0, jnp.log(x), 0.0) and expect 0.0 to be picked when $x\le0$. The forward result is indeed correct, but the gradient is still NaN.

Why: where can't block $0\times\infty$ in the backward pass¶

Let $y=\text{where}(c,\ a(x),\ b(x))$. jnp.where first computes both branches, then picks elementwise. In the backward pass, the upstream gradient $\bar y$ is split between the two branches according to the condition:

$$ \bar a = \text{where}(c,\ \bar y,\ 0),\qquad \bar b = \text{where}(c,\ 0,\ \bar y) $$

Then each branch keeps propagating back by the chain rule: $\bar x = \bar a\cdot a'(x) + \bar b\cdot b'(x)$.

When $c$ is false, $\bar a=0$. If $a'(x)$ is finite there, $0\times a'(x)=0$ and all is well. But if the dangerous branch has derivative $\pm\infty$ at that point (for example $\log'(0)=\frac10$) or NaN, you get

$$ 0\times\infty=\text{NaN}\quad(\text{per IEEE 754}) $$

NaN plus anything is still NaN, so $\bar x$ becomes NaN too. In other words, where only picks the forward value; it can't block the "$0\times\infty$" that the unselected branch brings in during the backward pass.

The "double where" trick¶

The root cause is that the dangerous branch was evaluated on the dangerous input. The fix: use a first where to replace the dangerous input with a safe value (any value that keeps the dangerous branch's derivative finite), then use a second where to pick the output:

x_safe = jnp.where(cond, x, SAFE)               # 1st where: the dangerous branch never sees the dangerous input
y = jnp.where(cond, danger_fn(x_safe), fallback)  # 2nd where: pick the output

Now the unselected branch's derivative $a'(\text{SAFE})$ is finite, and $0\times\text{finite}=0$. And when $c$ is false, the first where routes the gradient to the constant SAFE, so nothing flows back to $x$.

Common cases¶

First, take the smallest possible example apart by hand and watch $0\times\infty$ happen:

Practical tips

• The safe value has to keep the dangerous branch's derivative finite; a finite value alone isn't enough. Use 1.0 for sqrt/log/division and 0.0 for arcsin.
• If the masked-out branch has a finite derivative at the dangerous point, a single where is enough. For example, where(x>0, x**2, 0.) has gradient 0 at 0, with no NaN.
• When NaNs show up during training, you can turn on jax.config.update("jax_debug_nans", True) so JAX raises at the first primitive that produces a NaN, which pinpoints the problem quickly. It does miss some cases in this environment, though (for example after using vmap); see Advanced 08 §3.1 for usage and known limitations.
• For a fixed mathematical function, you can also supply a stable derivative directly with jax.custom_jvp (see below).

2. Numerical stability¶

• Vanishing / exploding gradients: in deep networks the chain rule multiplies many terms together, so gradients can shrink toward 0 or blow up.
• For example, $\frac{d}{dx}\log x=\frac1x$ is $10^{20}$ at $x=10^{-20}$, which is huge.

Ways to mitigate this:

• Gradient clipping;
• Use numerically stable implementations, e.g. jax.scipy.special.logsumexp(x) instead of jnp.log(jnp.sum(jnp.exp(x)));
• Turn on 64-bit precision when needed: jax.config.update("jax_enable_x64", True) (set at the start of the program, before creating any arrays).

🔍 Extra: custom derivatives with jax.custom_jvp¶

⚪ Optional: read this if automatic differentiation of some function is numerically unstable and you want to write the derivative rule yourself.

For some functions, the derivative you get by composing elementary operations is numerically unstable, even though the derivative itself has a simple, stable closed form. In that case you can use jax.custom_jvp to tell JAX directly how to compute the derivative. The classic example from the official docs is $\log(1+e^x)$ (softplus). Its derivative is $\sigma(x)=\frac{e^x}{1+e^x}$, but plain autodiff computes $\frac{\infty}{\infty}=\text{NaN}$ when $x$ is large.

You only define the JVP rule (forward mode), and JAX derives reverse mode automatically, so grad, jacrev and hessian all work. (There is also jax.custom_vjp, which lets you define the backward rule directly, for things like gradient clipping or saving memory.)

🔍 Extra: jax.lax.stop_gradient to keep part of a computation out of differentiation¶

⚪ Optional: read this if you need to cut off part of the gradient (target networks, straight-through estimators, etc.).

stop_gradient(x) returns x unchanged in the forward computation, but treats it as a constant (gradient 0) when differentiating. Typical uses: target networks in reinforcement learning, the stop-grad branch in contrastive learning, straight-through estimators, and so on.

3. Interaction with Python control flow¶

grad only differentiates the path the current input actually takes. If the function switches branches for different inputs, you only get the gradient of the current branch, which can be misleading near a branch boundary. With grad alone a Python if works fine, but under jit you need lax.cond / jnp.where instead (see 3.7).

4. External code and side effects¶

Inside a function you differentiate, use only JAX functions such as jnp, jax.lax, jax.scipy and jax.nn. External libraries behave in one of three ways under grad:

1. NumPy computations that don't depend on the input (like a constant array generated with np.random): they really are constants, and no gradient flows through them.
2. Passing a tracer to a NumPy function (like np.sin(x)): raises TracerArrayConversionError right away. Libraries like SciPy (not jax.scipy), Pandas and scikit-learn convert inputs to NumPy arrays internally, so they usually raise too.
3. A few NumPy functions (like np.sum) end up calling x.sum() through duck typing, so they happen to be differentiable. That relies on implementation details and isn't reliable; switch to jnp everywhere anyway.

Side effects such as printing, reading or writing files, and modifying globals break JAX's functional assumptions. Inside grad, print prints the tracer; inside jit, it runs only once, at trace time. To debug values, use jax.debug.print.

⚠️ Note: a NumPy call inside a differentiated function (like np.sum) is not silently treated as a constant with a zero gradient. np.sin(tracer) raises TracerArrayConversionError right away, while np.sum happens to differentiate correctly. Only NumPy constants that don't depend on the input are truly constants.

5. Non-scalar outputs¶

jax.grad only works for $f:\mathbb R^n\to\mathbb R$ (scalar output). If the function returns an array, it raises a TypeError. For vector-valued functions, use jax.jacfwd or jax.jacrev (for the full Jacobian), or jax.jvp and jax.vjp (for products of the Jacobian with a vector); see 3.3.

🔍 Extra: checking gradients with finite differences¶

⚪ Optional: read this if you wrote a custom derivative or a complex model and need to verify the gradients.

When you've written a custom_jvp, hand-written a backward rule, or built a fairly complex model, it's a good idea to check the gradients against numerical differentiation.

• Official tool: jax.test_util.check_grads(f, args, order, modes) uses central differences to check both forward (fwd) and reverse (rev) mode; order=2 also checks second derivatives. Note that you have to import jax.test_util explicitly. With only import jax, accessing jax.test_util raises AttributeError.
• By hand: central differences $\dfrac{f(x+h e_i)-f(x-h e_i)}{2h}$. The error is fairly large in float32, so use $h$ around $10^{-3}$ and loosen the tolerance accordingly.

3.9 Hands-on: computing gradients¶

This section uses a few exercises to practice jax.grad: starting from simple functions, then differentiating with respect to chosen arguments, higher-order derivatives, value_and_grad, and gradients of jax.numpy functions.

Setup: import jax, jax.numpy and numpy, turn on 64-bit precision if needed, then print the JAX version and the default backend.

Setup: how to turn on 64-bit precision¶

⚠️ Note: older material often does from jax.config import config followed by config.update("jax_enable_x64", True). JAX 0.12 has removed the jax.config submodule, so that raises ModuleNotFoundError: No module named 'jax.config'. Use jax.config.update("jax_enable_x64", True) instead, and put it at the start of the program, before creating any arrays (or use the environment variable JAX_ENABLE_X64=1).

To keep JAX's default float32 behavior for the rest of this chapter, the cells below only show the syntax and don't actually turn on 64-bit.

Exercise 1: gradient of a simple polynomial¶

$f(x)=3x^2+2x+5$, with analytic derivative $f'(x)=6x+2$. At $x=4$: $f'(4)=6\times4+2=26$.

Exercise 2: differentiating with respect to specific arguments¶

$g(x,y)=x^3y+2x^2$, with partial derivatives:

$$ \frac{\partial g}{\partial x}=3x^2y+4x,\qquad \frac{\partial g}{\partial y}=x^3 $$

With $x=2,\ y=3$: $\partial g/\partial x=3\times4\times3+8=44$, $\partial g/\partial y=8$.