SERIES: Learning RL and JAX in Public - from zero to DeepMind :)
Everyone talks about JAX's automatic differentiation and JIT compilation. Those are great. But vmap is the one that made me actually understand why research teams at places like DeepMind and Google Brain write everything in JAX. It is also the one that is hardest to explain with a one-liner, so let me take a proper shot at it.
The problem :
You have a function that works on one input. A single vector, a single example, a single data point.
def score(x):
return jnp.sum(x ** 2)
Now you have a batch of 1000 inputs and you need to run this on all of them.
The Python way:
results = [score(x) for x in batch] # slow loop
The PyTorch way: rewrite your function to handle batches explicitly, add batch dimensions, make sure everything broadcasts correctly. Tedious. Error-prone.
The JAX way: vmap !!
What vmap does :
vmap takes a function written for a single input and returns a new function that runs on a whole batch, in parallel, without you changing anything.
batch_score = jax.vmap(score)
results = batch_score(batch) # runs on all 1000 inputs at once
You wrote score thinking about one input. JAX handles the batching. Same results as the loop, but running as a single vectorised operation on the hardware.
I benchmarked this today. The vmap version was 32 times faster than a Python loop on CPU. On a GPU, that gap is even wider.
The in_axes argument :
This is the one thing about vmap that needs a second read. Some arguments you want to batch over, and some you want to keep fixed across the whole batch.
Example: you have a loss function with a shared weight vector w, but different inputs x and targets y for each example.
def loss(w, x, y):
prediction = jnp.dot(w, x)
return (prediction - y) ** 2
batch_loss = jax.vmap(loss, in_axes=(None, 0, 0))
in_axes=(None, 0, 0) means:
- None: do not batch over w. It is shared across all examples.
- 0: batch over the 0th axis of x. One x per example.
- 0: batch over the 0th axis of y. One y per example.
That is it. You tell vmap which arguments vary across the batch and which stay fixed.
The real JAX pattern: jit + vmap + grad together :
This is what DeepMind code looks like. One line, all three superpowers:
fast_batched_grad = jax.jit(jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0)))
Your loss function is now differentiable, batched across all examples simultaneously, and compiled for speed. In PyTorch this would take 10-15 lines and careful attention to tensor shapes.
I ran a full batched gradient descent loop today with this pattern. w started at [0, 0, 0] and converged to [2, 3, 5] in 115 steps with loss hitting exactly 0.0. I played with the iteration count to find that number and found the convergence satisfying to watch.
Why this matters for research :
When you are running experiments on large models, you often want to compute gradients for an entire batch of examples simultaneously. Without vmap you either loop (slow) or rewrite your functions to be batch-aware (messy and error-prone at scale).
With vmap you write clean single-example functions and compose them up to batch scale. Your code stays readable. Your functions are testable on one example at a time. And vmap handles the rest.
This is why JAX code in research papers is often remarkably compact compared to equivalent PyTorch. The functions stay simple. The composability does the scaling.
The summary I keep in my notes :
| What you want | How you get it |
|---|---|
| GPU acceleration | automatic |
| Gradients | jax.grad(fn) |
| Speed | jax.jit(fn) |
| Batch without a loop | jax.vmap(fn) |
| All three | jax.jit(jax.vmap(jax.grad(fn))) |
Next up: Day 4, which is where things get genuinely exciting. Q-learning from scratch. First real reinforcement learning algorithm. I built a gridworld environment in pure JAX and trained an agent to navigate it. See you there.
If you have been following along: which of the four superpowers surprised you most? For me it was vmap. I expected grad to be the revelation but vmap is the one that changed how I think about writing ML code.
All code from this series, organised by day, is on my GitHub: https://github.com/MadhumithaKolkar/jax-rl-lab
Happy learning everyone !
~ Madhumitha Kolkar (index_0)
Top comments (0)