Writing functions that handle batch dimensions is tedious and error-prone. You add extra dimensions, use einsum or manual broadcasting, and the code becomes hard to read. jax.vmap eliminates this entirely: write your function for a single example, and vmap automatically handles the batch.
import jax
import jax.numpy as jnp
# A function for ONE example
def predict_single(weights, x):
"""Linear prediction for a single input vector."""
return jnp.dot(weights, x)
weights = jnp.array([1.0, 2.0, 3.0])
single_x = jnp.array([0.5, 0.3, 0.8])
# Works for one example
print(predict_single(weights, single_x)) # 3.5
# Now handle a BATCH of examples — just use vmap!
batch_predict = jax.vmap(predict_single, in_axes=(None, 0))
batch_x = jnp.array([
[0.5, 0.3, 0.8],
[1.0, 0.0, 0.5],
[0.2, 0.4, 0.6],
])
# vmap automatically applies predict_single to each row
print(batch_predict(weights, batch_x)) # [3.5, 2.5, 2.8]
Without vmap, you'd have to write the batched version manually:
# Manual batching — error-prone and clutters the code
def predict_batch_manual(weights, X):
return X @ weights # Had to think about dimensions
# vmap version — no dimensional reasoning needed
predict_batch_vmap = jax.vmap(predict_single, in_axes=(None, 0))
# Both produce the same result, but vmap is cleaner
💡 Why This Matters
vmap isn't just syntactic sugar — it changes how you think about writing numerical code. Instead of always reasoning about batch dimensions, you write the simplest possible function (for one example) and let vmap handle batching. This makes code easier to read, test, and debug. It's also more composable: the single-example version works naturally with grad, and then vmap(grad(f)) gives you batched gradients.