Skip to content

← JAX Foundations step 12 of 25

Medium Primitives

jit + grad Composition

Jump to the task

Compile the derivative function

jax.grad(loss) builds a function that computes a scalar loss's gradient. jax.jit can compile that function, just as it can compile the loss itself. Apply both transformations to the function before passing in its data:

import jax
import jax.numpy as jnp

def loss(w, x):
    return jnp.sum((w * x) ** 2)

compiled_grad = jax.jit(jax.grad(loss, argnums=0))
compiled_grad(2.0, jnp.array([1.0, 3.0]))  # 40
compiled_grad(3.0, jnp.array([1.0, 3.0]))  # 60

Here argnums=0 selects w. The second call can reuse the compiled program: the values changed, but the argument shapes and dtypes did not. Reuse the same transformed function across calls instead of defining a new loss and JIT wrapper inside every invocation.

What compilation changes

The first call normally pays tracing and compilation costs. Repeated compatible calls can reuse that work. Different shapes, dtypes or static argument values can require another compilation. A speedup is workload-dependent; this tiny example is meant to teach composition, not demonstrate a particular performance gain. When timing device computations, wait for the returned result with .block_until_ready().

jax.grad(jax.jit(loss)) is also valid for this loss and computes the same mathematical derivative. The placement changes the compilation boundary, not which argument is differentiated. An outer JIT is a convenient way to compile the whole gradient calculation. If you also need the loss value, jax.jit(jax.value_and_grad(loss)) returns a pair containing the value and gradient rather than just the gradient.

This diagram illustrates cache reuse for repeated shapes. The lengths 32 and 64 are illustrative, not requirements of the exercise:

Loading visualization…

Problem

Implement jit_grad_quadratic(w, b, x) using jax.jit(jax.grad(...)). Define the scalar loss sum((w*x + b)**2) and differentiate with respect to w only. The inputs w and b are real floating-point scalars; x is a one-dimensional real floating-point JAX array. Return a scalar:

gradient = 2 * sum((w*x + b) * x)

For w=0, b=1 and x=[1, 2], return 6. For w=1, b=2 and x=[3], return 30. These examples include the bias contribution even though b is not the argument being differentiated.

Define the loss and compiled gradient at module scope, then call the compiled function from the entrypoint. Pass the loss function itself to grad; grad(loss(w, b, x)) instead passes a scalar where a callable is required. Output tests check the derivative values, not whether your implementation uses JIT or reuses its compilation cache.