Topolearnvmap removes error-prone batch-dimension code and gives per-example gradients almost for free.
jax.vmap(f) vectorizes f over a leading batch axis. You write f for a single example; vmap produces a function that accepts a stack of examples and computes all of them at once, as efficiently as hand-written batched code, without a Python loop.
in_axes says which axis of each argument to map over; None means "broadcast this argument unchanged":
batched = jax.vmap(predict, in_axes=(None, 0)) # params shared, x batchedout_axes controls where the batch axis appears in the output (default 0). Arguments can be pytrees, with in_axes given per subtree.
Because transformations compose, vmap combines with the others:
jax.jit(jax.vmap(f))compiles the batched function.jax.vmap(jax.grad(loss), in_axes=(None, 0, 0))gives per-example gradients, which is awkward in most frameworks.- Nested
vmapmaps over several axes, e.g. pairwise distances between two sets of points.
vmap traces the function like jit does, so the body should not branch in Python on array values. All mapped arguments must have the same size along their mapped axis, otherwise you get a clear ValueError.
Think of vmap as "write it for one, run it for many". It is often simpler and faster than reshaping arrays to fit a batched formula by hand.
Example
import jax
import jax.numpy as jnp
def predict(w, x): # one example: x has shape (3,)
return jnp.tanh(jnp.dot(w, x))
w = jnp.array([0.5, -1.0, 2.0])
xs = jnp.ones((4, 3)) # a batch of 4 examples
batched = jax.vmap(predict, in_axes=(None, 0))
print(batched(w, xs).shape) # (4,)
def loss(w, x, y):
return (predict(w, x) - y) ** 2
ys = jnp.zeros(4)
per_example = jax.vmap(jax.grad(loss), in_axes=(None, 0, 0))(w, xs, ys)
print(per_example.shape) # (4, 3): one gradient per example
# pairwise distances with nested vmap
a, b = jnp.ones((2, 3)), jnp.zeros((5, 3))
dist = jax.vmap(jax.vmap(lambda p, q: jnp.linalg.norm(p - q), (None, 0)), (0, None))
print(dist(a, b).shape) # (2, 5)Remember
- Write the function for one example;
vmapadds the batch axis in_axes=Nonebroadcasts an argument such as shared parametersvmap(grad(...))yields per-example gradients
Watch out for
- Mapping arguments with different sizes along the mapped axis, which raises a ValueError.