What grad Does: Functions In, Gradient Functions Out

In calculus, taking the derivative of a function gives you a new function. JAX's jax.grad does exactly this: it takes a Python function that returns a scalar, and returns a new Python function that computes the gradient.

import jax
import jax.numpy as jnp

# A scalar-valued function
def f(x):
    return x ** 2

# grad(f) returns a NEW function that computes df/dx
df = jax.grad(f)

print(f(3.0))    # 9.0
print(df(3.0))   # 6.0 (derivative of x^2 is 2x, evaluated at x=3)

# df is a regular Python function — you can call it, JIT it, etc.
print(df(5.0))   # 10.0

This is fundamentally different from PyTorch's approach. In PyTorch, you compute a loss, call loss.backward(), and then read param.grad. The gradient is attached to the tensor. In JAX, grad is a function transformation — it takes a function and returns a function. No state is mutated anywhere.

import jax
import jax.numpy as jnp

# Multi-variable function
def loss(params, x, y):
    w, b = params
    pred = jnp.dot(x, w) + b
    return jnp.mean((pred - y) ** 2)

# grad differentiates with respect to the FIRST argument by default
grad_fn = jax.grad(loss)

# Call it: returns gradient with same structure as params
w = jnp.array([1.0, 2.0])
b = jnp.array(0.0)
params = (w, b)
x = jnp.array([[1.0, 0.5], [0.3, 0.8]])
y = jnp.array([1.0, 0.5])

grads = grad_fn(params, x, y)
print(type(grads))       # tuple — same structure as params!
print(grads[0].shape)    # (2,) — gradient of w
print(grads[1].shape)    # () — gradient of b