Differentiating with respect to different arguments#

By default jax.grad differentiates with respect to the first argument. The argnums parameter selects other arguments, or several at once. Here’s a logistic regression model, where we might want gradients with respect to the weights W, the bias b, or both:

def sigmoid(x):
  return 0.5 * (jnp.tanh(x / 2) + 1)

# Outputs probability of a label being true.
def predict(W, b, inputs):
  return sigmoid(jnp.dot(inputs, W) + b)

# A toy dataset, and some parameter values.
inputs = jnp.array([[0.52, 1.12,  0.77],
                    [0.88, -1.08, 0.15],
                    [0.52, 0.06, -1.30],
                    [0.74, -2.49, 1.39]])
targets = jnp.array([True, True, False, True])
W = jnp.array([0.1, 0.4, -0.3])
b = 0.5

# Training loss is the negative log-likelihood of the training examples.
def loss(W, b):
  preds = predict(W, b, inputs)
  label_probs = preds * targets + (1 - preds) * (1 - targets)
  return -jnp.sum(jnp.log(label_probs))

# Differentiate `loss` with respect to the first positional argument:
W_grad = jax.grad(loss, argnums=0)(W, b)
print(f'{W_grad=}')

# Since argnums=0 is the default, this does the same thing:
W_grad = jax.grad(loss)(W, b)
print(f'{W_grad=}')

# But you can choose different values too, and drop the keyword:
b_grad = jax.grad(loss, 1)(W, b)
print(f'{b_grad=}')

# Including tuple values:
W_grad, b_grad = jax.grad(loss, (0, 1))(W, b)
print(f'{W_grad=}')
print(f'{b_grad=}')

W_grad=Array([-0.72034794,  1.9374751 , -2.226889  ], dtype=float32)
W_grad=Array([-0.72034794,  1.9374751 , -2.226889  ], dtype=float32)
b_grad=Array(-0.7629071, dtype=float32, weak_type=True)
W_grad=Array([-0.72034794,  1.9374751 , -2.226889  ], dtype=float32)
b_grad=Array(-0.7629071, dtype=float32, weak_type=True)

Real models don’t keep their parameters in separate positional arguments. Instead they use nested containers like dictionaries. jax.grad handles those natively; that’s the subject of the next page, Pytrees.

value_and_grad and auxiliary outputs#

You usually want the loss value as well as its gradient, for example to log training progress. jax.value_and_grad() computes both in one pass:

loss_value, Wb_grad = jax.value_and_grad(loss, (0, 1))(W, b)
print(loss_value)

3.5066924

And sometimes a function naturally computes intermediate results worth returning alongside the scalar being differentiated. Have the function return a (scalar_output, aux_data) pair, and pass has_aux=True:

def loss_and_preds(W, b):
  preds = predict(W, b, inputs)
  label_probs = preds * targets + (1 - preds) * (1 - targets)
  return -jnp.sum(jnp.log(label_probs)), preds

W_grad, preds = jax.grad(loss_and_preds, has_aux=True)(W, b)
print(preds)

[0.68330455 0.5277215  0.7243215  0.30174544]

Checking derivatives numerically#

Derivatives are easy to check against finite differences:

eps = 1e-4
b_grad_numerical = (loss(W, b + eps / 2.) - loss(W, b - eps / 2.)) / eps
print('b_grad_numerical', b_grad_numerical)
print('b_grad_autodiff', jax.grad(loss, 1)(W, b))

b_grad_numerical -0.76293945
b_grad_autodiff -0.7629071

JAX ships a convenience that does this automatically, to any order:

from jax.test_util import check_grads

check_grads(loss, (W, b), order=2)  # check up to 2nd order derivatives

jax.grad is one entry point into a much deeper autodiff system built on the fundamental jax.jvp() and jax.vjp() transformations. Those are covered in the advanced autodiff docs; see JAX 301: advanced autodiff and extending JAX.

Choosing which axes to map with in_axes and out_axes#

By default, vmap maps over the leading axis of every input. The in_axes and out_axes arguments override this. For example, if your data has the batch as the second axis:

auto_batch_convolve_v2 = jax.vmap(convolve, in_axes=1, out_axes=1)

xst = jnp.transpose(xs)
wst = jnp.transpose(ws)

auto_batch_convolve_v2(xst, wst)

Array([[11., 11.],
       [20., 20.],
       [29., 29.]], dtype=float32)

An in_axes entry of None means “don’t map this argument”: the same value is passed to every call. Here we convolve a batch of xs against one shared w:

batch_convolve_v3 = jax.vmap(convolve, in_axes=(0, None))

batch_convolve_v3(xs, w)

Array([[11., 20., 29.],
       [11., 20., 29.]], dtype=float32)

Composing vmap#

Like all JAX transformations, vmap composes. Nesting vmap gives concise expressions for “all pairs” computations:

def dist(x, y):
  return jnp.sqrt(jnp.sum((x - y) ** 2))

def all_pairs(f):
  return jax.vmap(jax.vmap(f, in_axes=(None, 0)), in_axes=(0, None))

points = jnp.array([[0., 0.], [1., 0.], [0., 2.]])
all_pairs(dist)(points, points)

Array([[0.      , 1.      , 2.      ],
       [1.      , 0.      , 2.236068],
       [2.      , 2.236068, 0.      ]], dtype=float32)

Composing vmap with grad answers questions that are awkward to express otherwise. For example: the gradient of our logistic regression loss for each example separately (rather than summed over the batch). Write the per-example loss, differentiate it, then vectorize the derivative:

def example_loss(W, b, x, y):
  pred = predict(W, b, x)
  label_prob = pred * y + (1 - pred) * (1 - y)
  return -jnp.log(label_prob)

per_example_grads = jax.vmap(jax.grad(example_loss, (0, 1)),
                             in_axes=(None, None, 0, 0))
W_grads, b_grads = per_example_grads(W, b, inputs, targets)
print(W_grads)  # one gradient per example
print(b_grads)

[[-0.16468163 -0.3546989  -0.24385549]
 [-0.41560507  0.5100608  -0.07084177]
 [ 0.3766471   0.04345929 -0.94161785]
 [-0.5167084   1.7386539  -0.97057384]]
[-0.31669545 -0.47227848  0.7243214  -0.6982546 ]

Each transformation did one conceptually simple job, and composition did the rest. This is the characteristic JAX pattern: write the mathematically natural, single-example function, and build everything else out of transformations.

Consequence 1: transformations require traceable functions#

JAX transformations only work on operations that the tracing machinery intercepts. The usual sufficient condition for traceability is functional purity: outputs depend only on inputs (arguments and closed-over values); outputs are produced by applying JAX operations; and no side effects occur.

Purity is valuable even without a tracing implementation: it makes code easier to reason about, and easier for the compiler to optimize, parallelize, and scale. It also gives the transformations simple, clear meanings. A pure function denotes a mathematical function, and the transformations are mathematical operators on it: jax.grad(f) means \(\nabla f\); jax.vmap(f) means “\(f\) applied to each element,” with no question of which side effects happen or in what order; jax.jit(f) can promise to return exactly what f returns, while caching and optimizing freely.

Tracing then turns this from good advice into a working requirement. Side effects in your function, like built-in Python print calls, happen at trace time, not when the transformed computation runs. Our vmap example above already showed this: the batch had three elements, but print ran only once, because vmap traces the function a single time, transforming each operation as it’s intercepted. Under jax.jit, the effect is even sharper: traces are cached, so a side effect might happen on the first call and then never again. If you want to print runtime values from transformed code, there’s a purpose-built tool: jax.debug.print(); see Debugging runtime values.

Reading mutable state has the same problem in reverse: a global value is likely baked in at trace time, so later updates to it are silently ignored by transformed code.

For how to work with state in traceable code, see Stateful computations.

Consequence 2: traced code can’t always specialize on data values#

How much does a tracer know about the value it stands in for? That depends on the transformation. vmap’s tracers actually know quite a lot: they carry the whole batch of values along as the trace proceeds. We can even peek at the batch mid-trace, through the tracer’s .val attribute:

# warning: `.val` is unsupported internals, so don't rely on it in real code!
jax.vmap(lambda x: print(x.val))(jnp.arange(3.0))

[0. 1. 2.]

So the values may well be present. What traced code can’t do is specialize on them. The operations applied must work for every element of the batch, so the traced code can’t take one particular control-flow branch based on one particular element’s value. Each element might want a different branch, and the trace can only record one:

def absolute(x):
  if x > 0:      # needs one answer, but x stands for a whole batch of values
    return x
  else:
    return -x

jax.vmap(absolute)(jnp.arange(-2.0, 3.0))

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
This VmapTracer with object id 125716124378384 was created on line:
  /tmp/ipykernel_1497/4198118709.py:2:5 (absolute)
See https://docs.jax.dev/en/latest/101/errors.html#jax.errors.TracerBoolConversionError

The error message links to the fix: express data-dependent choices as array operations, like jax.numpy.where(), which compute a per-element answer instead of forcing a single branch:

def absolute(x):
  return jnp.where(x > 0, x, -x)

jax.vmap(absolute)(jnp.arange(-2.0, 3.0))

Array([ 2.,  1., -0.,  1.,  2.], dtype=float32)

(For data-dependent loops and more elaborate control flow, JAX provides structured control-flow operations like jax.lax.cond and jax.lax.scan; see Control flow and logical operators with jit.)

jax.jit is the extreme case: its tracers carry no values at all, only the JAX type, because its recorded program must serve every value of that type. The same constraint follows, for an even stronger reason: there are no values to consult in the first place.

On the other hand, anything that depends only on JAX types works freely during tracing, because those are ordinary Python values at trace time. Python for loops over a fixed range, if statements on shapes, and shape arithmetic are all fine. Our convolve function above used a Python loop whose bounds came from len(x): that loop unrolls during tracing, and the recorded operations are as if we’d written the unrolled version by hand. This is also why keeping import numpy as np around is useful: np operations on shapes execute immediately at trace time, separating “computations on static values” from “computations being traced” (jnp).

grad works with Python control flow#

If a vmap tracer stands for a whole batch of values, and a jit tracer stands for any value of the right JAX type, jax.grad’s tracers occupy the opposite pole: applied on its own, jax.grad evaluates your function with tracers that carry exactly one concrete value alongside the derivative bookkeeping. With a single value, data-dependent Python control flow is unambiguous, so if you’re just using grad, you can apply data-dependent if, while, recursion, or anything else:

def f(x):
  if x < 3:
    return 3. * x ** 2
  else:
    return -4. * x

print(jax.grad(f)(2.0))  # differentiates the 3x² branch: 6x at x=2
print(jax.grad(f)(4.0))  # differentiates the -4x branch

12.0
-4.0

Each call differentiates the branch actually taken, giving exactly the piecewise derivative you’d write by hand. Even loops whose trip count depends on the data are fine. Here’s differentiating through Newton’s method for square roots, where the number of iterations depends on the input:

def sqrt_newton(a):
  x = a
  while abs(x * x - a) > 1e-6:   # data-dependent loop
    x = 0.5 * (x + a / x)
  return x

print(sqrt_newton(2.0))
print(jax.grad(sqrt_newton)(2.0))  # 1/(2√2) ≈ 0.3536

1.4142135623746899
0.35355338

This flexibility is part of JAX’s lineage: JAX’s autodiff grew out of Autograd, whose whole point was differentiating ordinary, idiomatic Python and NumPy code, including branches, loops, and closures.

The constraints return the moment you compose with a transformation that traces abstractly: under jax.jit(jax.grad(f)) or jax.vmap(jax.grad(f)), the if above fails again, and needs rewriting with jnp.where, lax.cond, and friends (Control flow and logical operators with jit). That trade is this documentation’s split in miniature: grad by itself maximizes what you can express, and it’s compiling for speed that asks you to make control flow explicit.

Where jit fits in#

jax.jit() uses this same tracing machinery, but instead of replacing each intercepted operation on the fly, it records them all and hands the recording to the XLA compiler to produce fast fused machine code, caching the compiled result keyed on the JAX types of its inputs. In use, it’s most often applied as a decorator:

@jax.jit          # equivalent to sum_sq = jax.jit(sum_sq)
def sum_sq(x):
  return jnp.sum(x ** 2)

sum_sq(jnp.arange(3.0))

Array(5., dtype=float32)

The first call with a given set of input JAX types pays for tracing and compilation; later calls skip straight to the compiled code. The rest of the jit story is about performance (compilation, caching and retracing, static arguments, asynchronous dispatch), and it’s where the performance and scaling docs start; see Just-in-time compilation.