Automatic differentiation with grad

Gradients drive optimization and training; jax.grad gives exact derivatives of ordinary numeric code.

jax.grad(f) returns a new function that computes the gradient of f with respect to its first argument. f must return a real scalar; for vector outputs you need Jacobians (an Advanced lesson).

df = jax.grad(f)           # derivative w.r.t. first arg
df_dy = jax.grad(f, argnums=1)

Because grad is just a function transformation, it composes: jax.grad(jax.grad(f)) is the second derivative. The argument you differentiate with respect to must be floating point; passing an integer raises an error, a common surprise with f(3) instead of f(3.0).

Two variants you will use constantly:

jax.value_and_grad(f) returns (f(x), grad) in one pass, so a training step does not evaluate the loss twice.
has_aux=True lets f return (loss, aux); only loss is differentiated and aux (metrics, predictions) is passed through.

JAX computes gradients by reverse-mode autodiff: it traces the operations f performs and runs them backwards, which costs a small constant multiple of evaluating f no matter how many inputs there are. Python control flow (if, for) is fine under plain grad because tracing follows the branch taken for the concrete value; it becomes a restriction once you add jit.

The argument can be any pytree, such as a dict of weight arrays; the gradient then has the same structure.

Example

import jax
import jax.numpy as jnp

def f(x):
    return x ** 3 + 2.0 * x

df = jax.grad(f)
print(df(2.0))                  # 14.0   (3x^2 + 2)
print(jax.grad(df)(2.0))        # 12.0   (6x)

def loss(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

w = jnp.array([1.0, 0.0])
x = jnp.array([[1.0, 2.0], [3.0, 4.0]])
y = jnp.array([1.0, 2.0])
val, g = jax.value_and_grad(loss)(w, x, y)
print(val, g)                   # 0.5 [3. 4.]

Remember

jax.grad(f) returns a gradient function w.r.t. the first argument (change with argnums)
The differentiated function must return a scalar and take float inputs
Use value_and_grad and has_aux in training steps

Watch out for

Calling the gradient function with an int like df(2), which raises a TypeError about integer inputs.

Check yourself

For f(x, y) = x * y**2, what does jax.grad(f, argnums=1)(3.0, 2.0) return?

Official docs: docs.jax.dev/en/latest/101/transformations.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks