Chapter 6¶

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
results = []
for x in batch:
 results.append(f(x))
result = jnp.array(results) # [2. 5. 10. 17. 26.]

# Method 2
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
# B
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
# Default (out_axes=0)
# out_axes=1

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
# Y
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
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
# Second vmap
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
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
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
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
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'
# 'b'

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