mx.grad and the Function-Transform Family

Autograd by composition, not by tape

PyTorch's autograd works by recording a tape of operations as you do them, then walking the tape backward. MLX (and JAX) take a different route — they expose transformations that take a function and return a new function. mx.grad(f) doesn't differentiate f's tape after the fact; it returns a new function that, when called, computes the gradient.

That sounds like a small distinction. It is not. It means your gradients are first-class composable values — you can wrap them in other transforms, JIT-compile them, vmap them across a batch. None of that requires a special "training mode" or a tape that stays live.

The three you'll use most

mx.grad(f) — given a scalar-valued function, returns a function that computes the gradient with respect to its first argument. Use it when you want gradients but don't care about the loss value.

mx.value_and_grad(f) — same as mx.grad, but returns both the loss value and the gradient in one call. Use this in training loops; you almost always want to log the loss.

mx.vmap(f) — turns a function-of-one-example into a function-of-a-batch. Composes naturally with grad. Use it when you want efficient batched computation without manually broadcasting.

Composability is the point

Because each transform takes a function and returns a function, you can stack them. mx.vmap(mx.grad(loss_fn)) gives you the per-example gradient across a batch. mx.grad(mx.grad(f)) gives you the second derivative. mx.compile(mx.value_and_grad(loss_fn)) JIT-compiles the gradient computation. Other frameworks make these tricks possible; MLX makes them obvious.

nn.value_and_grad — the model-aware variant

When you have an nn.Module with parameters, you usually want gradients with respect to the model's parameters, not just the function's first argument. The nn module exposes nn.value_and_grad(model, loss_fn) which is the canonical pattern for training. Lesson 6 will use it.

Code

mx.grad — the one-liner·python

import mlx.core as mx

def square(x):
    return x ** 2

grad_sq = mx.grad(square)
print(grad_sq(mx.array(3.0)))   # → array(6, dtype=float32)   d/dx(x^2) = 2x = 6

value_and_grad — the training-loop primitive·python

import mlx.core as mx

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

vg = mx.value_and_grad(loss)

w0 = mx.array(1.5)
xs = mx.array([1.0, 2.0, 3.0, 4.0])
ys = mx.array([2.0, 4.0, 6.0, 8.0])      # true relation y = 2x

v, g = vg(w0, xs, ys)
print('loss value:', float(v))            # → 1.875
print('gradient w.r.t. w:', float(g))     # → -7.5  (you'd subtract this in SGD)

vmap — batched function calls·python

import mlx.core as mx

def f(x):
    return x ** 2 + 1

# Vectorize f over the leading axis
batched = mx.vmap(f)

xb = mx.array([1.0, 2.0, 3.0, 4.0])
print(batched(xb))   # → array([2, 5, 10, 17], dtype=float32)

# Composes with grad: per-example gradient across a batch
batched_grad = mx.vmap(mx.grad(f))
print(batched_grad(xb))   # → array([2, 4, 6, 8], dtype=float32)   (= 2x for each)

Exercise

Pick a tiny scalar function of two arguments (e.g., def f(x, y): return (x - y) ** 2 + x). Use mx.grad(f, argnums=0) and mx.grad(f, argnums=1) to compute partial derivatives with respect to x and y. Verify the answers by hand. Then compose mx.vmap over a batch of (x, y) pairs. The point of the exercise is to feel that argnums + composition give you full control without writing a custom backward pass.