Vmap

Write a function for one example and call jax.vmap on it, and JAX does not generate a Python loop that calls your function N times. It rewrites the traced program itself: every primitive inside gets replaced by its batched form, a matmul becomes a batched matmul, an elementwise op just grows an extra axis, and the whole thing comes out as one fused program that XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → compiles once. Nothing in the result dispatches per example; the batch dimension is built into the primitives themselves.

It's tempting to picture vmap as sugar for a loop, something that saves you typing for i in range(n). That mental model gets the performance story backward. A real loop dispatches N separate programs; vmap produces a single program that happens to carry an extra axis through it, and the two only resemble each other from the outside.

The contract you can rely on is exact, and it reads as an equation: vmap(f)(xs)[i] == f(xs[i]), for every i, up to the order floating-point additions happen to associate in. Whatever f computes for one example, vmap(f) computes for every example in the batch independently, and the result is indistinguishable except in the narrow sense that summing floats in a different order can move the last bit.

vmap rewrites the program, not the loop.

vmap needs to know, for every argument, where its batch axis lives, and in_axes is how you say it. Pass an integer and that's the axis holding the batch dimension for that argument. Pass None and the argument is broadcast instead: every call sees the whole thing, unbatched, with no axis to strip out first. When your arguments are pytrees rather than bare arrays, in_axes can itself be a matching pytree of specs, one entry per leaf.

out_axes plays the same role for the output side, choosing where the newly introduced batch dimension lands in whatever f returns. And because vmap is just another transformation, you can nest it: wrap a vmapped function in another vmap and the axes stack, reading from the inside out. The innermost vmap is the one that ran first, on the function you actually wrote.

One requirement vmap enforces and won't relax: at least one argument has to carry a mapped axis. Set every entry of in_axes to None and JAX raises, because without a single batch dimension anywhere in the inputs, there's nothing telling it how large the output's new axis should even be.

import jax
import jax.numpy as jnp

def dot(a, b):
    return a @ b

a = jnp.ones((32, 8))
b = jnp.ones(8)
print(jax.vmap(dot, in_axes=(0, None))(a, b).shape)   # (32,): b broadcast

pair = jax.vmap(jax.vmap(dot, in_axes=(None, 0)), in_axes=(0, None))
print(pair(a, a).shape)   # (32, 32): every pair of rows, two nested maps

The composition worth knowing by heart is jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)). Read it in order: grad turns the per-example loss into a per-example gradient function, and vmap then runs that gradient function once per example in the batch, with the parameters broadcast (their in_axes entry is None) and the data batched across the other two. What comes out is one gradient per training example, computed at roughly the cost of a single batched matmul rather than a genuine Python loop over gradients. This is the pattern behind per-example gradient clipping and differential-privacy accounting, anything that needs the gradient before it gets averaged away.

The Jacobian builders from the last chapter are built the same way underneath. jacfwd is vmap wrapped around jvp; jacrev is vmap wrapped around vjp. Batching the single-direction cost across every direction at once is, quite literally, what a Jacobian is.

The same trick scales up to whole models. Stack N sets of parameters into one pytree, an extra leading axis on every leaf, and vmap a training step over that stacked pytree, and you're training N independent models inside a single compiled program: an ensemble without N separate training loops competing for the compiler's attention.

import jax
import jax.numpy as jnp

def loss(w, x, y):
    return (x @ w - y) ** 2

w = jnp.ones(3)
xs = jnp.ones((32, 3))
ys = jnp.zeros(32)
per_sample = jax.vmap(jax.grad(loss), in_axes=(None, 0, 0))(w, xs, ys)
print(per_sample.shape)   # (32, 3): one gradient per example