Effortless Batching with jax.vmap

Introduction

Welcome to the second lesson of our course "Advanced JAX: Transformations for Speed & Scale"! I'm delighted to continue this journey with you as we dive deeper into JAX's powerful transformation capabilities. In our previous lesson, you mastered JAX's functional approach to randomness, learning how to create, split, and manage PRNG keys for reproducible computations. That foundation of explicit state management — rather than relying on hidden global state — perfectly sets the stage for today's topic. In this lesson, we'll explore automatic vectorization with jax.vmap, one of JAX's most elegant and powerful transformations. Just as you learned to explicitly manage randomness through keys, vmap allows you to explicitly control how computations are batched across data, transforming functions that operate on single examples into functions that efficiently process entire batches. This transformation is fundamental to high-performance scientific computing and machine learning, where we routinely need to apply the same operation to thousands or millions of data points. By the end of this lesson, you'll understand how to use jax.vmap to automatically vectorize your functions, control batching behavior with the in_axes and out_axes parameters, and achieve significant performance improvements over manual Python loops. You'll also see how vmap seamlessly integrates with other JAX transformations like jit, creating a powerful toolkit for scalable computations.

Understanding the Batching Challenge

Before we dive into jax.vmap, let's understand the common challenge it solves. In scientific computing and machine learning, we frequently write functions that operate on individual data points — a single vector, matrix, or more complex structure. However, in practice, we almost always need to apply these functions to batches of data: hundreds of images, thousands of training examples, or millions of simulation parameters. The traditional approach to this challenge involves manually writing loops or reshaping data, which often leads to several problems: Performance Issues : Python loops are notoriously slow, especially when processing large amounts of numerical data. Each iteration involves Python interpreter overhead. Code Complexity : Manual batching often requires careful indexing, reshaping, and accumulation of results, making code harder to read and debug. Memory Inefficiency : Naive implementations might create unnecessary intermediate arrays or fail to leverage vectorized operations. Transformation Incompatibility : Handwritten loops don't compose well with JAX transformations like jit or automatic differentiation. Consider a simple scenario: you have a function that computes the mean and standard deviation of a single vector, but you need to apply it to a batch of vectors. The manual approach might involve a Python for loop that calls your function repeatedly and collects results. While functional, this approach sacrifices both performance and elegance. This is precisely where automatic vectorization becomes invaluable. The idea is to take a function designed for single data points and automatically transform it into a function that efficiently processes entire batches, without requiring you to rewrite the original logic.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Python

import jax
import jax.numpy as jnp

def process_single_vector(vector):
    """Computes mean and standard deviation of a single vector."""
    return jnp.mean(vector), jnp.std(vector)

# Create a batch of vectors (5 vectors, each of size 10)
key = jax.random.key(101)
batch_of_vectors = jax.random.normal(key, shape=(5, 10))
print(f"Batch of vectors (shape {batch_of_vectors.shape}):\n{batch_of_vectors}")

import jax
import jax.numpy as jnp

def process_single_vector(vector):
    """Computes mean and standard deviation of a single vector."""
    return jnp.mean(vector), jnp.std(vector)

# Create a batch of vectors (5 vectors, each of size 10)
key = jax.random.key(101)
batch_of_vectors = jax.random.normal(key, shape=(5, 10))
print(f"Batch of vectors (shape {batch_of_vectors.shape}):\n{batch_of_vectors}")

text

Batch of vectors (shape (5, 10)):
[[-0.8633027   1.564201    0.6821926   0.9258969   2.501565    0.7184641
  -0.25084117  1.392589   -1.1157503   0.3450869 ]
  ... ]

Batch of vectors (shape (5, 10)):
[[-0.8633027   1.564201    0.6821926   0.9258969   2.501565    0.7184641
  -0.25084117  1.392589   -1.1157503   0.3450869 ]
  ... ]

Python

# Using vmap to apply process_single_vector to the batch
# in_axes=0 means map over the 0-th axis of batch_of_vectors
# out_axes=0 means the outputs will also be batched along axis 0
batched_process_fn = jax.vmap(process_single_vector, in_axes=0, out_axes=0)

means, stds = batched_process_fn(batch_of_vectors)
print("\nUsing jax.vmap:")
print(f"Means for each vector: {means}")
print(f"Standard deviations for each vector: {stds}")

# Using vmap to apply process_single_vector to the batch
# in_axes=0 means map over the 0-th axis of batch_of_vectors
# out_axes=0 means the outputs will also be batched along axis 0
batched_process_fn = jax.vmap(process_single_vector, in_axes=0, out_axes=0)

means, stds = batched_process_fn(batch_of_vectors)
print("\nUsing jax.vmap:")
print(f"Means for each vector: {means}")
print(f"Standard deviations for each vector: {stds}")

text

Using jax.vmap:
Means for each vector: [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ]
Standard deviations for each vector: [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ]

Using jax.vmap:
Means for each vector: [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ]
Standard deviations for each vector: [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ]

Python

import time # Make sure time is imported

# For comparison: Python loop (less efficient)
print("\nUsing Python loop (for comparison):")
loop_means = []
loop_stds = []
start_time_loop = time.perf_counter()
for i in range(batch_of_vectors.shape[0]):
    mean, std = process_single_vector(batch_of_vectors[i])
    loop_means.append(mean)
    loop_stds.append(std)
end_time_loop = time.perf_counter()

# JIT the vmapped function for a fair speed comparison
jit_batched_process_fn = jax.jit(batched_process_fn)
# Warm-up JIT: run once so compilation time isn't included in benchmark
_, _ = jit_batched_process_fn(batch_of_vectors) 

start_time_vmap = time.perf_counter()
jit_means, jit_stds = jit_batched_process_fn(batch_of_vectors)
# Ensure computation is complete before stopping the timer
jit_means.block_until_ready() 
jit_stds.block_until_ready()
end_time_vmap = time.perf_counter()

print(f"Loop means: {jnp.array(loop_means)}") # Convert list to JAX array for consistent display
print(f"Loop stds: {jnp.array(loop_stds)}")
print(f"Time for Python loop: {end_time_loop - start_time_loop:.6f}s")
print(f"Time for JITted vmap: {end_time_vmap - start_time_vmap:.6f}s")

import time # Make sure time is imported

# For comparison: Python loop (less efficient)
print("\nUsing Python loop (for comparison):")
loop_means = []
loop_stds = []
start_time_loop = time.perf_counter()
for i in range(batch_of_vectors.shape[0]):
    mean, std = process_single_vector(batch_of_vectors[i])
    loop_means.append(mean)
    loop_stds.append(std)
end_time_loop = time.perf_counter()

# JIT the vmapped function for a fair speed comparison
jit_batched_process_fn = jax.jit(batched_process_fn)
# Warm-up JIT: run once so compilation time isn't included in benchmark
_, _ = jit_batched_process_fn(batch_of_vectors) 

start_time_vmap = time.perf_counter()
jit_means, jit_stds = jit_batched_process_fn(batch_of_vectors)
# Ensure computation is complete before stopping the timer
jit_means.block_until_ready() 
jit_stds.block_until_ready()
end_time_vmap = time.perf_counter()

print(f"Loop means: {jnp.array(loop_means)}") # Convert list to JAX array for consistent display
print(f"Loop stds: {jnp.array(loop_stds)}")
print(f"Time for Python loop: {end_time_loop - start_time_loop:.6f}s")
print(f"Time for JITted vmap: {end_time_vmap - start_time_vmap:.6f}s")

text

Using Python loop (for comparison):
Loop means: [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ]
Loop stds: [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ]
Time for Python loop: 0.061947s
Time for JITted vmap: 0.000026s

Using Python loop (for comparison):
Loop means: [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ]
Loop stds: [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ]
Time for Python loop: 0.061947s
Time for JITted vmap: 0.000026s

Python

def process_vector_and_scalar(vector, scalar_offset):
    """Adds a scalar offset to the mean of a vector."""
    return jnp.mean(vector) + scalar_offset

# vmap with multiple arguments: some batched, some not (broadcasted)
print("\nUsing jax.vmap with multiple arguments (batched and broadcasted):")
scalar_value = 100.0
# Map over batch_of_vectors (axis 0 for the first arg), 
# broadcast scalar_value (None for the second arg)
batched_mixed_fn = jax.vmap(process_vector_and_scalar, in_axes=(0, None), out_axes=0)

results_mixed = batched_mixed_fn(batch_of_vectors, scalar_value)
print(f"Results (mean + scalar_offset) for each vector:\n{results_mixed}")

def process_vector_and_scalar(vector, scalar_offset):
    """Adds a scalar offset to the mean of a vector."""
    return jnp.mean(vector) + scalar_offset

# vmap with multiple arguments: some batched, some not (broadcasted)
print("\nUsing jax.vmap with multiple arguments (batched and broadcasted):")
scalar_value = 100.0
# Map over batch_of_vectors (axis 0 for the first arg), 
# broadcast scalar_value (None for the second arg)
batched_mixed_fn = jax.vmap(process_vector_and_scalar, in_axes=(0, None), out_axes=0)

results_mixed = batched_mixed_fn(batch_of_vectors, scalar_value)
print(f"Results (mean + scalar_offset) for each vector:\n{results_mixed}")

text

Using jax.vmap with multiple arguments (batched and broadcasted):
Results (mean + scalar_offset) for each vector:
[100.59001 100.3225  100.01243 100.02004 100.25158]

Using jax.vmap with multiple arguments (batched and broadcasted):
Results (mean + scalar_offset) for each vector:
[100.59001 100.3225  100.01243 100.02004 100.25158]

Python

# Example: batching over a different axis
# Transpose the original batch: shape (10, 5) - 10 features, 5 examples
vectors_alt_batch = batch_of_vectors.T 
print(f"\nAlternative batch format (shape {vectors_alt_batch.shape})")

# Map over axis 1 (columns) of vectors_alt_batch, 
# stack outputs along axis 0 (default for out_axes with multiple outputs)
batched_alt_fn = jax.vmap(process_single_vector, in_axes=1, out_axes=0)
means_alt, stds_alt = batched_alt_fn(vectors_alt_batch)

print(f"Means (alt batching): {means_alt}, shape: {means_alt.shape}")
print(f"Stds (alt batching): {stds_alt}, shape: {stds_alt.shape}")

# Example: batching over a different axis
# Transpose the original batch: shape (10, 5) - 10 features, 5 examples
vectors_alt_batch = batch_of_vectors.T 
print(f"\nAlternative batch format (shape {vectors_alt_batch.shape})")

# Map over axis 1 (columns) of vectors_alt_batch, 
# stack outputs along axis 0 (default for out_axes with multiple outputs)
batched_alt_fn = jax.vmap(process_single_vector, in_axes=1, out_axes=0)
means_alt, stds_alt = batched_alt_fn(vectors_alt_batch)

print(f"Means (alt batching): {means_alt}, shape: {means_alt.shape}")
print(f"Stds (alt batching): {stds_alt}, shape: {stds_alt.shape}")

text

Alternative batch format (shape (10, 5))
Means (alt batching): [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ], shape: (5,)
Stds (alt batching): [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ], shape: (5,)

Alternative batch format (shape (10, 5))
Means (alt batching): [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ], shape: (5,)
Stds (alt batching): [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ], shape: (5,)

Python

# Using vmap as a decorator
@jax.vmap
def vectorized_normalize(vector):
    """Normalizes a vector to unit length (decorator approach)."""
    norm = jnp.linalg.norm(vector)
    return vector / norm

# Test the decorated function
normalized_vectors = vectorized_normalize(batch_of_vectors)
print(f"Normalized vectors shape: {normalized_vectors.shape}")
print(f"Norms of normalized vectors: {jnp.linalg.norm(normalized_vectors, axis=1)}")

# Using vmap as a decorator
@jax.vmap
def vectorized_normalize(vector):
    """Normalizes a vector to unit length (decorator approach)."""
    norm = jnp.linalg.norm(vector)
    return vector / norm

# Test the decorated function
normalized_vectors = vectorized_normalize(batch_of_vectors)
print(f"Normalized vectors shape: {normalized_vectors.shape}")
print(f"Norms of normalized vectors: {jnp.linalg.norm(normalized_vectors, axis=1)}")

Python

from functools import partial

# Pre-configure vmap for specific batching patterns
@partial(jax.vmap, in_axes=(0, None), out_axes=0)
def scale_vectors(vectors, scale_factor):
    """Scales vectors by a common factor."""
    return vectors * scale_factor

# Use the pre-configured decorator
scaled_vectors = scale_vectors(batch_of_vectors, 2.0)
print(f"Scaled vectors shape: {scaled_vectors.shape}")
print(f"First original vector: {batch_of_vectors[0]}")
print(f"First scaled vector: {scaled_vectors[0]}")

from functools import partial

# Pre-configure vmap for specific batching patterns
@partial(jax.vmap, in_axes=(0, None), out_axes=0)
def scale_vectors(vectors, scale_factor):
    """Scales vectors by a common factor."""
    return vectors * scale_factor

# Use the pre-configured decorator
scaled_vectors = scale_vectors(batch_of_vectors, 2.0)
print(f"Scaled vectors shape: {scaled_vectors.shape}")
print(f"First original vector: {batch_of_vectors[0]}")
print(f"First scaled vector: {scaled_vectors[0]}")

Python

import jax
import jax.numpy as jnp

def process_single_vector(vector):
    """Computes mean and standard deviation of a single vector."""
    return jnp.mean(vector), jnp.std(vector)

# Create a batch of vectors (5 vectors, each of size 10)
key = jax.random.key(101)
batch_of_vectors = jax.random.normal(key, shape=(5, 10))
print(f"Batch of vectors (shape {batch_of_vectors.shape}):\n{batch_of_vectors}")

text

Batch of vectors (shape (5, 10)):
[[-0.8633027   1.564201    0.6821926   0.9258969   2.501565    0.7184641
  -0.25084117  1.392589   -1.1157503   0.3450869 ]
  ... ]

Python

# Using vmap to apply process_single_vector to the batch
# in_axes=0 means map over the 0-th axis of batch_of_vectors
# out_axes=0 means the outputs will also be batched along axis 0
batched_process_fn = jax.vmap(process_single_vector, in_axes=0, out_axes=0)

means, stds = batched_process_fn(batch_of_vectors)
print("\nUsing jax.vmap:")
print(f"Means for each vector: {means}")
print(f"Standard deviations for each vector: {stds}")

text

Using jax.vmap:
Means for each vector: [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ]
Standard deviations for each vector: [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ]

Python

import time # Make sure time is imported

# For comparison: Python loop (less efficient)
print("\nUsing Python loop (for comparison):")
loop_means = []
loop_stds = []
start_time_loop = time.perf_counter()
for i in range(batch_of_vectors.shape[0]):
    mean, std = process_single_vector(batch_of_vectors[i])
    loop_means.append(mean)
    loop_stds.append(std)
end_time_loop = time.perf_counter()

# JIT the vmapped function for a fair speed comparison
jit_batched_process_fn = jax.jit(batched_process_fn)
# Warm-up JIT: run once so compilation time isn't included in benchmark
_, _ = jit_batched_process_fn(batch_of_vectors) 

start_time_vmap = time.perf_counter()
jit_means, jit_stds = jit_batched_process_fn(batch_of_vectors)
# Ensure computation is complete before stopping the timer
jit_means.block_until_ready() 
jit_stds.block_until_ready()
end_time_vmap = time.perf_counter()

print(f"Loop means: {jnp.array(loop_means)}") # Convert list to JAX array for consistent display
print(f"Loop stds: {jnp.array(loop_stds)}")
print(f"Time for Python loop: {end_time_loop - start_time_loop:.6f}s")
print(f"Time for JITted vmap: {end_time_vmap - start_time_vmap:.6f}s")

text

Using Python loop (for comparison):
Loop means: [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ]
Loop stds: [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ]
Time for Python loop: 0.061947s
Time for JITted vmap: 0.000026s

Python

def process_vector_and_scalar(vector, scalar_offset):
    """Adds a scalar offset to the mean of a vector."""
    return jnp.mean(vector) + scalar_offset

# vmap with multiple arguments: some batched, some not (broadcasted)
print("\nUsing jax.vmap with multiple arguments (batched and broadcasted):")
scalar_value = 100.0
# Map over batch_of_vectors (axis 0 for the first arg), 
# broadcast scalar_value (None for the second arg)
batched_mixed_fn = jax.vmap(process_vector_and_scalar, in_axes=(0, None), out_axes=0)

results_mixed = batched_mixed_fn(batch_of_vectors, scalar_value)
print(f"Results (mean + scalar_offset) for each vector:\n{results_mixed}")

text

Using jax.vmap with multiple arguments (batched and broadcasted):
Results (mean + scalar_offset) for each vector:
[100.59001 100.3225  100.01243 100.02004 100.25158]

Python

# Example: batching over a different axis
# Transpose the original batch: shape (10, 5) - 10 features, 5 examples
vectors_alt_batch = batch_of_vectors.T 
print(f"\nAlternative batch format (shape {vectors_alt_batch.shape})")

# Map over axis 1 (columns) of vectors_alt_batch, 
# stack outputs along axis 0 (default for out_axes with multiple outputs)
batched_alt_fn = jax.vmap(process_single_vector, in_axes=1, out_axes=0)
means_alt, stds_alt = batched_alt_fn(vectors_alt_batch)

print(f"Means (alt batching): {means_alt}, shape: {means_alt.shape}")
print(f"Stds (alt batching): {stds_alt}, shape: {stds_alt.shape}")

text

Alternative batch format (shape (10, 5))
Means (alt batching): [0.5900101  0.3225023  0.01243191 0.02004608 0.251578  ], shape: (5,)
Stds (alt batching): [1.0570618  0.97786814 1.2293593  0.6301977  0.9018042 ], shape: (5,)

Python

# Using vmap as a decorator
@jax.vmap
def vectorized_normalize(vector):
    """Normalizes a vector to unit length (decorator approach)."""
    norm = jnp.linalg.norm(vector)
    return vector / norm

# Test the decorated function
normalized_vectors = vectorized_normalize(batch_of_vectors)
print(f"Normalized vectors shape: {normalized_vectors.shape}")
print(f"Norms of normalized vectors: {jnp.linalg.norm(normalized_vectors, axis=1)}")

Python

from functools import partial

# Pre-configure vmap for specific batching patterns
@partial(jax.vmap, in_axes=(0, None), out_axes=0)
def scale_vectors(vectors, scale_factor):
    """Scales vectors by a common factor."""
    return vectors * scale_factor

# Use the pre-configured decorator
scaled_vectors = scale_vectors(batch_of_vectors, 2.0)
print(f"Scaled vectors shape: {scaled_vectors.shape}")
print(f"First original vector: {batch_of_vectors[0]}")
print(f"First scaled vector: {scaled_vectors[0]}")