Chapter 6: Vectorization - jax.vmap¶
Overview¶
jax.vmap automatically vectorizes functions - transforms a function that works on single elements to work on batches.
The Vectorization Problem¶
Manual Batching (Slow)¶
import jax.numpy as jnp
def f(x):
return x**2 + 1
# Process batch manually
batch = jnp.array([1., 2., 3., 4., 5.])
# Method 1: Python loop (very slow)
results = []
for x in batch:
results.append(f(x))
result = jnp.array(results) # [2. 5. 10. 17. 26.]
# Method 2: Vectorize manually (tedious)
def f_batch(batch):
return batch**2 + 1
result = f_batch(batch) # Works but had to rewrite f
Automatic Vectorization (vmap)¶
import jax
import jax.numpy as jnp
def f(x):
return x**2 + 1
# Automatic vectorization
f_vmap = jax.vmap(f)
batch = jnp.array([1., 2., 3., 4., 5.])
result = f_vmap(batch) # [2. 5. 10. 17. 26.] - automatic!
Key insight: Write function for single element, vmap handles batching.
Basic Usage¶
Single Batch Dimension¶
import jax
import jax.numpy as jnp
def f(x):
return jnp.sin(x) + jnp.cos(x)
# Vectorize along first axis
f_vmap = jax.vmap(f)
# Single value
print(f(1.0)) # ~1.38
# Batch (first dimension mapped)
batch = jnp.array([0.0, 1.0, 2.0])
print(f_vmap(batch)) # [1.0, 1.38, -0.15]
Multiple Arguments¶
import jax
import jax.numpy as jnp
def f(x, y):
return x**2 + y**3
f_vmap = jax.vmap(f)
x = jnp.array([1., 2., 3.])
y = jnp.array([4., 5., 6.])
result = f_vmap(x, y) # Element-wise operations
print(result) # [65., 133., 225.]
2D Arrays (Batch Matrix Operations)¶
import jax
import jax.numpy as jnp
def matrix_multiply(A, B):
return A @ B # Matrix multiply
# vmap over batch dimension
batch_multiply = jax.vmap(matrix_multiply)
# A: (3, 2, 3) = 3 matrices of shape (2, 3)
# B: (3, 3, 4) = 3 matrices of shape (3, 4)
A = jnp.ones((3, 2, 3))
B = jnp.ones((3, 3, 4))
result = batch_multiply(A, B)
print(result.shape) # (3, 2, 4) - 3 result matrices
Vectorization Axes: in_axes and out_axes¶
in_axes: Which Axes to Batch¶
import jax
import jax.numpy as jnp
def f(x):
return x**2
# Batch along axis 0 (default)
f_vmap_0 = jax.vmap(f, in_axes=0)
# Batch along axis 1
f_vmap_1 = jax.vmap(f, in_axes=1)
# Example
x = jnp.array([[1., 2., 3.],
[4., 5., 6.]])
print(f_vmap_0(x)) # Apply f to each row
# [[1., 4., 9.],
# [16., 25., 36.]]
print(f_vmap_1(x)) # Apply f to each column
# [[1., 16.],
# [4., 25.],
# [9., 36.]]
Multiple Arguments with Different Axes¶
import jax
import jax.numpy as jnp
def f(x, y):
return x * y
# vmap x along axis 0, y along axis 1
f_vmap = jax.vmap(f, in_axes=(0, 1))
x = jnp.array([1., 2., 3.]) # Shape (3,)
y = jnp.array([[1., 2., 3.],
[4., 5., 6.],
[7., 8., 9.]]) # Shape (3, 3)
result = f_vmap(x, y)
print(result.shape) # (3, 3)
No Batching for Some Arguments¶
import jax
import jax.numpy as jnp
def f(x, const):
return x * const
# Batch x (axis 0), don't batch const (None)
f_vmap = jax.vmap(f, in_axes=(0, None))
x = jnp.array([1., 2., 3.])
const = 5.0
result = f_vmap(x, const) # [5., 10., 15.]
out_axes: Output Batch Dimension¶
import jax
import jax.numpy as jnp
def f(x):
return jnp.array([x, x**2, x**3]) # Returns (3,) for each input
# Output shape: (3,) per input, 5 inputs
# Default (out_axes=0): (5, 3) - batch dimension first
# out_axes=1: (3, 5) - output dimension first
f_vmap_0 = jax.vmap(f, in_axes=0, out_axes=0)
f_vmap_1 = jax.vmap(f, in_axes=0, out_axes=1)
x = jnp.array([1., 2., 3., 4., 5.])
result_0 = f_vmap_0(x)
print(result_0.shape) # (5, 3)
result_1 = f_vmap_1(x)
print(result_1.shape) # (3, 5)
Composing with Other Transformations¶
vmap + grad¶
import jax
import jax.numpy as jnp
from jax import grad, vmap
def f(x):
return x**3
# Gradient of each element in batch
grad_f = grad(f)
vmap_grad_f = vmap(grad_f)
x = jnp.array([1., 2., 3.])
result = vmap_grad_f(x) # [3., 12., 27.]
vmap + jit¶
import jax
import jax.numpy as jnp
from jax import jit, vmap
def f(x):
return jnp.sin(x)**2 + jnp.cos(x)**3
# Vectorize and compile
f_fast = jit(vmap(f))
x = jnp.array([1., 2., 3., 4., 5.])
result = f_fast(x) # Fast and batched!
vmap(grad(...))¶
import jax
import jax.numpy as jnp
from jax import grad, vmap
def loss(w, x, y):
return (w * x - y)**2
# Gradient for each sample in batch
w = jnp.array([1., 2., 3.])
X = jnp.array([[1., 2., 3.],
[4., 5., 6.],
[7., 8., 9.]])
y = jnp.array([1., 2., 3.])
# Gradient of loss w.r.t. w for each sample
grad_loss = vmap(grad(loss, argnums=0), in_axes=(None, 0, None))
grads = grad_loss(w, X, y)
print(grads.shape) # (3, 3) - 3 samples, 3 parameters
Practical Examples¶
Batch Dot Products¶
import jax
import jax.numpy as jnp
def dot(x, y):
return jnp.dot(x, y)
batch_dot = jax.vmap(dot)
# X: (32, 10) - 32 samples, 10-dim vectors
# Y: (32, 10) - 32 samples, 10-dim vectors
X = jnp.ones((32, 10))
Y = jnp.ones((32, 10))
results = batch_dot(X, Y) # (32,) - dot product for each sample
Batch Neural Network Forward Pass¶
import jax
import jax.numpy as jnp
from jax import vmap
def forward(params, x):
"""Forward pass for single input"""
w1, b1, w2, b2 = params
h = jnp.maximum(0, x @ w1 + b1) # ReLU
return h @ w2 + b2
# Vectorized forward pass
forward_batch = vmap(forward, in_axes=(None, 0))
# Parameters (shared)
params = (
jnp.ones((784, 128)),
jnp.zeros(128),
jnp.ones((128, 10)),
jnp.zeros(10)
)
# Batch of 32 samples
X_batch = jnp.ones((32, 784))
# Compute outputs for all samples
outputs = forward_batch(params, X_batch)
print(outputs.shape) # (32, 10)
Batch Jacobian¶
import jax
import jax.numpy as jnp
from jax import vmap, jacobian
def f(x):
return jnp.array([x[0]**2, x[1]**3])
# Jacobian for each sample in batch
jac = jacobian(f)
batch_jac = vmap(jac)
X = jnp.array([[1., 2.], [3., 4.], [5., 6.]])
jacobians = batch_jac(X)
print(jacobians.shape) # (3, 2, 2)
Nested vmap¶
Multiple Batch Dimensions¶
import jax
import jax.numpy as jnp
def f(x):
return x**2
# Apply vmap twice
f_vmap_2 = jax.vmap(jax.vmap(f))
# 2D array: (3, 4)
x = jnp.arange(12).reshape((3, 4)).astype(float)
result = f_vmap_2(x) # (3, 4) - applied to each element
print(result.shape) # (3, 4)
Different Axes per Level¶
import jax
import jax.numpy as jnp
def f(x, y):
return x * y
# First vmap: batch x along axis 0
# Second vmap: batch y along axis 0
f_vmap_2 = jax.vmap(jax.vmap(f, in_axes=(None, 0)), in_axes=(0, None))
x = jnp.array([1., 2., 3.]) # (3,)
y = jnp.array([[1., 2.],
[3., 4.]]) # (2, 2)
# Result: (3, 2, 2)
result = f_vmap_2(x, y)
Broadcasting vs vmap¶
Broadcasting¶
import jax.numpy as jnp
x = jnp.array([1., 2., 3.]) # Shape (3,)
y = 5.0
# Broadcasting: scalar expands
result = x + y # [6., 7., 8.]
vmap¶
import jax
import jax.numpy as jnp
def f(x):
# x is scalar
return x**2
# vmap applies f to each element
f_vmap = jax.vmap(f)
x = jnp.array([1., 2., 3.])
result = f_vmap(x) # [1., 4., 9.]
Difference: Broadcasting is element-wise on arrays. vmap applies function to each element of batch dimension.
Common Patterns¶
Pattern 1: Batch Processing in Models¶
import jax
import jax.numpy as jnp
from jax import vmap
class Model:
def __init__(self, params):
self.params = params
def forward(self, x):
# Single sample
return x @ self.params['w']
def predict_batch(self, X):
# Batch: apply forward to each sample
return vmap(self.forward)(X)
Pattern 2: Gradient Per Sample¶
import jax
import jax.numpy as jnp
from jax import grad, vmap
def loss(w, x, y):
return (w @ x - y)**2
# Get gradient for each sample
grad_per_sample = vmap(grad(loss, argnums=0), in_axes=(None, 0, 0))
w = jnp.array([1., 2., 3.])
X = jnp.ones((32, 3)) # 32 samples
y = jnp.ones(32)
grads = grad_per_sample(w, X, y) # (32, 3)
Pattern 3: Batch Normalization¶
import jax
import jax.numpy as jnp
from jax import vmap
def normalize_single(x):
return (x - jnp.mean(x)) / jnp.std(x)
# Normalize each sample
normalize_batch = vmap(normalize_single)
X = jnp.ones((32, 784))
X_normalized = normalize_batch(X)
Performance Characteristics¶
Benchmarking¶
import jax
import jax.numpy as jnp
import time
def f(x):
return jnp.sin(x)**2 + jnp.cos(x)**3
# Method 1: Python loop
x_list = [jnp.array(i) for i in range(10000)]
start = time.time()
results = [f(x) for x in x_list]
print(f"Python loop: {time.time() - start:.3f}s")
# Method 2: vmap
x_array = jnp.arange(10000.)
f_vmap = jax.vmap(f)
start = time.time()
result = f_vmap(x_array)
print(f"vmap: {time.time() - start:.3f}s") # Much faster!
Common Gotchas¶
Gotcha 1: Shape Mismatch¶
import jax
import jax.numpy as jnp
def f(x):
return x @ jnp.ones(10) # Expects shape (10,)
f_vmap = jax.vmap(f)
# ❌ Wrong shape
x = jnp.ones((5, 8)) # (5, 8) - doesn't work!
# ✅ Correct shape
x = jnp.ones((5, 10)) # (5, 10) - batch of (10,) vectors
result = f_vmap(x)
Gotcha 2: Wrong Axis¶
import jax
import jax.numpy as jnp
def f(x):
return jnp.sum(x)
f_vmap_0 = jax.vmap(f, in_axes=0)
f_vmap_1 = jax.vmap(f, in_axes=1)
x = jnp.array([[1., 2., 3.],
[4., 5., 6.]])
print(f_vmap_0(x)) # [3., 15.] - sum each row
print(f_vmap_1(x)) # [5., 7., 9.] - sum each column
Gotcha 3: Unexpected Shapes with Functions Returning Structured Data¶
import jax
import jax.numpy as jnp
def f(x):
return {'a': x**2, 'b': x**3}
f_vmap = jax.vmap(f)
x = jnp.array([1., 2., 3.])
result = f_vmap(x)
# result = {'a': array([1., 4., 9.]),
# 'b': array([1., 8., 27.])}
Summary¶
VMAP QUICK REFERENCE:
Basic:
f_vmap = jax.vmap(f)
result = f_vmap(batch)
Custom axes:
f_vmap = jax.vmap(f, in_axes=1, out_axes=0)
No batch for some args:
f_vmap = jax.vmap(f, in_axes=(0, None, 0))
Nested vmap:
f_vmap_2 = jax.vmap(jax.vmap(f))
Compose:
f_fast = jax.jit(jax.vmap(jax.grad(loss)))
Next Steps¶
- 04 Parallelization Jax.Pmap - Multi-device parallelization
- 08 Composing Transformations - Combine vmap, grad, jit
Checkpoint: Vectorize a function that computes y = 3x² + 2x + 1 for a batch of 1000 values
Last Updated: 2026-08-09