Nested vmap and Advanced Patterns

You can nest vmap calls to vectorize over multiple dimensions. Each vmap adds one batch dimension.

import jax
import jax.numpy as jnp

def pairwise_distance(a, b):
    """Euclidean distance between two vectors."""
    return jnp.sqrt(jnp.sum((a - b) ** 2))

# Vectorize over pairs: distance between corresponding rows
batched_distance = jax.vmap(pairwise_distance)

A = jnp.array([[1.0, 0.0], [3.0, 0.0], [0.0, 4.0]])
B = jnp.array([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])

print(batched_distance(A, B))  # [1.0, 3.0, 4.0]

# Double vmap: pairwise distance MATRIX (every a against every b)
# Outer vmap: over rows of A
# Inner vmap: over rows of B
distance_matrix = jax.vmap(jax.vmap(pairwise_distance, in_axes=(None, 0)), in_axes=(0, None))

D = distance_matrix(A, B)
print(D.shape)  # (3, 3)
print(D)
# [[1.    0.    0.   ]   distances from A[0] to each B
#  [3.    0.    0.   ]   distances from A[1] to each B
#  [4.    0.    0.   ]]  distances from A[2] to each B

vmap + jit: Maximum Performance

import jax
import jax.numpy as jnp

def process_single(params, x):
    """Complex per-example computation."""
    h = jnp.tanh(x @ params['w1'] + params['b1'])
    return jnp.sum(h ** 2)

# Stack all transformations: vectorize, then compile
process_batch = jax.jit(jax.vmap(process_single, in_axes=(None, 0)))

# Create params as a dict (pytree)
key = jax.random.PRNGKey(0)
k1, k2 = jax.random.split(key)
params = {
    'w1': jax.random.normal(k1, (4, 8)),
    'b1': jnp.zeros(8),
}

X = jax.random.normal(k2, (128, 4))  # batch of 128

result = process_batch(params, X)
print(result.shape)  # (128,) — one scalar per example

Data Augmentation with vmap

import jax
import jax.numpy as jnp

def augment_single(key, image):
    """Random augmentation for a single image."""
    k1, k2 = jax.random.split(key)
    # Random brightness
    brightness = jax.random.uniform(k1, minval=0.8, maxval=1.2)
    # Random noise
    noise = 0.05 * jax.random.normal(k2, image.shape)
    return jnp.clip(image * brightness + noise, 0.0, 1.0)

# Vectorize over a batch of (key, image) pairs
augment_batch = jax.jit(jax.vmap(augment_single))

# Each image gets its own random key for independent augmentation
master_key = jax.random.PRNGKey(42)
batch_size = 32
keys = jax.random.split(master_key, batch_size)
images = jax.random.uniform(jax.random.PRNGKey(0), (batch_size, 28, 28))

augmented = augment_batch(keys, images)
print(augmented.shape)  # (32, 28, 28)