Chapter 4: Automatic Vectorization with vmap
Personal study notes on the course "Getting Started with JAX" · Environment: Python 3.14 / JAX 0.12.0.1 / CPU
Original course page: Chapter 4 overview
In numerical computing and machine learning, you often need to apply the same function to a whole batch of data, which is called "batching". Writing explicit loops, or rewriting a function specifically for batches, tends to be verbose and forces you to manage array dimensions by hand. jax.vmap is JAX's transformation for automatic vectorization: you write a function for "a single data point", and it can apply that function to a batch (or several batches) of data, usually without changing the function's logic. You can think of it as automatically adding a "batch dimension" to the computation.
Learning goals (the 6 listed in the course)¶
- Understand what vectorization means and why it improves performance;
- Use
jax.vmapto vectorize functions that take one or more arguments; - Control which axes are mapped with
in_axesandout_axes; - Handle more complex batching structures with nested
vmap; - Combine
vmapwithjitandgrad; - Know the performance and memory considerations when using
vmap.
Sections¶
| Section | Original page |
|---|---|
| 4.1 The concept of vectorization | concept-of-vectorization |
4.2 Introducing jax.vmap | introducing-jax-vmap |
4.3 Mapping over specific arguments (in_axes, out_axes) | vmap-axes-arguments |
| 4.4 Handling multiple batched arguments | vmap-multiple-arguments |
4.5 Nested vmap | nesting-vmap |
4.6 Combining vmap with jit and grad | combining-vmap-jit-grad |
4.7 Performance considerations for vmap | vmap-performance |
| 4.8 Hands-on: vectorizing functions | hands-on-vectorizing-functions |
| 4.9 Chapter 4 Quiz | quiz |
How to read this
- Each section follows the same structure: course explanation (key points reorganized) → course code (comments in English, using the recommended JAX 0.12 style throughout) → 🔍 Extra.
- Common pitfalls are marked with
⚠️ Note. When code is involved, the outdated or wrong version goes in a ❌ cell, followed right away by the ✅ recommended version, and the conclusion is checked withassert. These notes are collected in the "Version notes" at the end of the chapter.- "Vectorization" here always means vectorization, not quantization (the course pages occasionally mix up the two terms).
- All code and conclusions were actually run and verified on JAX 0.12.0.1 (CPU).
The explanations are rewritten around the course's key points. The code examples keep the course's variable names and values where possible, and all of them have been run.
# Common imports and helpers for this chapter
import time
import jax
import jax.numpy as jnp
from jax import lax
import numpy as np
print("JAX version:", jax.__version__)
print("Devices:", jax.devices())
def bench(fn, *args, n=20):
# Warm up once (triggers compilation), then time n runs and average, in milliseconds
fn(*args).block_until_ready()
t0 = time.perf_counter()
for _ in range(n):
out = fn(*args)
out.block_until_ready()
return (time.perf_counter() - t0) / n * 1e3
def jaxpr_str(closed_jaxpr):
# Turn a jaxpr into plain text without terminal color codes, so we can check its contents with assert
return closed_jaxpr.pretty_print(use_color=False)
4.1 The concept of vectorization¶
The problem. In scientific computing and machine learning you often apply the same operation to many different inputs: processing a batch of images, running a simulation many times with different parameters, computing each training example's contribution to the loss.
In Python, the most direct approach is a for loop. Say you have a function predict(image) that handles one image; you call it once per image. That's easy to read, but it is often a performance bottleneck:
- Every iteration pays interpreter overhead;
- The operations inside the loop handle one example at a time, so they can't make full use of the parallelism of GPUs/TPUs;
- Python is interpreted, so loops are much slower than optimized compiled code.
What is vectorization? Vectorization turns "an operation on a scalar (or a small, fixed-size input)" into an operation that acts on a whole array or a whole batch at once. At the level of high-level code, instead of processing elements one by one in a loop, it processes all the data points in one go.
As a formula: applying the same function $f$ to a batch of inputs $x_1,\dots,x_N$,
$$ y_i = f(x_i),\ i = 1,\dots,N \qquad\Longleftrightarrow\qquad Y = F(X),\quad X = \begin{bmatrix} x_1 \\ \vdots \\ x_N \end{bmatrix} $$
The left side is "one at a time" (a loop); the right side is "all at once": stack the $N$ examples into an array with one extra dimension and compute everything with the batched version $F$.
Why is vectorization faster? Three main reasons:
- Less interpreter overhead: one high-level instruction triggers an operation over the whole array, with no loop counter to maintain in Python and no fetching elements one by one.
- Optimized low-level routines: NumPy and JAX are built on highly optimized routines (such as BLAS/LAPACK) tuned for CPU, GPU and TPU architectures.
- Hardware parallelism: modern CPUs have SIMD instructions that apply the same operation to several values at once, and GPUs/TPUs have thousands of compute units. Vectorized code uses this hardware naturally.
Besides being fast, vectorized code is usually more concise: numbers + 10 says what you mean more directly than an explicit loop.
Vectorization in JAX. NumPy functions like np.add and np.sin are already vectorized, but you have to organize your inputs and code according to broadcasting rules. JAX's jax.vmap goes a step further: you write the logic for a single data point (like predict(image)), and vmap automatically turns it into a function that handles a whole batch efficiently. You don't rewrite any batching logic or manage any loops.
Course code¶
# Course code 1: process a batch of images one at a time (predict is only illustrative; a placeholder implementation is added so it runs)
def predict(image):
return jnp.mean(image) # Placeholder: return one score per image
def predict_batch_loop(images):
results = []
for image in images:
results.append(predict(image))
return results # You could also stack these into one array with jnp.stack
images = jnp.ones((3, 8, 8)) # 3 "images" of size 8x8
print(predict_batch_loop(images))
# Course code 2: a plain Python loop
numbers = [1, 2, 3, 4, 5]
results = []
for x in numbers:
results.append(x + 10)
print(results) # [11, 12, 13, 14, 15]
# Course code 3: the vectorized version (NumPy)
import numpy as np
numbers = np.array([1, 2, 3, 4, 5])
results = numbers + 10 # Adds to the whole array at once, with no explicit Python loop
print(results) # [11 12 13 14 15]
assert results.tolist() == [11, 12, 13, 14, 15]
🔍 Extra: measuring the speed gap between a loop and vectorization¶
⚪ Optional: read this if you want to see the measured speed difference between loops and vectorization.
Let's measure it: 1000 vectors of dimension 64, each dotted with the same vector.
key = jax.random.key(0)
k1, k2 = jax.random.split(key)
xs = jax.random.normal(k1, (1000, 64)) # 1000 examples, each a 64-dimensional vector
w = jax.random.normal(k2, (64,))
def dot_loop(xs, w): # Python loop: every iteration dispatches a device operation
return jnp.stack([jnp.dot(x, w) for x in xs])
def dot_manual(xs, w): # Manual vectorization: one matrix-vector product
return xs @ w
assert jnp.allclose(dot_loop(xs, w), dot_manual(xs, w), atol=1e-5)
print("Python loop: %.2f ms" % bench(dot_loop, xs, w, n=3))
print("Manual vectorization: %.3f ms" % bench(dot_manual, xs, w))
🔍 Extra: where manual vectorization most often goes wrong¶
🟢 Must-read
NumPy-style vectorization asks you to organize your code around broadcasting rules, and that is exactly where bugs creep in most easily:
- Add one dimension to the input and the same function means something else. For example
jnp.dot(a, b): for two 1-D arrays it's an inner product, but for two 2-D arrays it's a matrix product. Feeding a batch directly into single-example code often raises no error but gives the wrong result. - Other common "silent bugs":
x.sum()orjnp.linalg.norm(x)were meant to work per example but end up covering the whole batch.x.Talso means different things for 2-D and 3-D arrays.
In the example below, the single-example function normalizes a vector to unit length, $x / \lVert x \rVert$. Passing the whole batch directly raises no error, but jnp.linalg.norm on a 2-D array computes the Frobenius norm of the whole matrix, so every row gets divided by the same number.
def normalize(x): # Single example: x has shape (3,); returns a unit vector
return x / jnp.linalg.norm(x)
X = jnp.array([[3.0, 4.0, 0.0],
[1.0, 0.0, 0.0]]) # 2 examples; in the correct result every row has norm 1
wrong = normalize(X) # Pass the whole batch directly: no error, but the result is wrong
right = jax.vmap(normalize)(X) # vmap: each example is normalized on its own
print("Whole batch passed directly:\n", wrong, "\n row norms:", jnp.linalg.norm(wrong, axis=1))
print("vmap:\n", right, "\n row norms:", jnp.linalg.norm(right, axis=1))
assert not jnp.allclose(jnp.linalg.norm(wrong, axis=1), 1.0) # Whole batch passed directly: wrong result
assert jnp.allclose(jnp.linalg.norm(right, axis=1), 1.0) # vmap: every row is a unit vector
4.2 Introducing jax.vmap¶
Applying a function written for a single data point to a whole batch is an extremely common need. Python loops or NumPy's built-in vectorization handle some of these cases, but explicitly managing dimensions and loops buries the core logic. That gets worse as the function becomes more complex or has several batched inputs.
jax.vmap (vectorizing map) is a function transformation designed for exactly this. It automatically adds a batch dimension to a function without touching the function's internal logic. You only think and write in terms of "a single example", and vmap turns that into a function that handles a whole batch efficiently.
The course demonstrates this with a "sum of squares of a vector" function: first call it on a single vector, then process a 4×3 batch (each row is a vector to be processed independently) with a manual loop and with vmap.
Course code¶
import jax
import jax.numpy as jnp
# A function that processes a single vector (expects a 1-D array)
def sum_of_squares(vector):
print(f"Running sum_of_squares for a vector of shape: {vector.shape}")
return jnp.sum(vector**2)
single_vector = jnp.array([1., 2., 3.])
result_single = sum_of_squares(single_vector)
print(f"Result for single vector: {result_single}") # 1^2 + 2^2 + 3^2 = 14.0
# A batch of vectors: 4 rows, each a 3-dimensional vector
batch_of_vectors = jnp.array([
[1., 2., 3.],
[4., 5., 6.],
[7., 8., 9.],
[0., 1., 0.],
])
print(f"Batch shape: {batch_of_vectors.shape}") # (4, 3)
# Manual loop (for demonstration only; not recommended in JAX)
manual_results = []
for i in range(batch_of_vectors.shape[0]):
vector = batch_of_vectors[i]
manual_results.append(sum_of_squares(vector))
manual_output = jnp.stack(manual_results)
print(f"Manual loop output: {manual_output}") # The print inside the function runs 4 times
# Create the vectorized version with vmap
vectorized_sum_of_squares = jax.vmap(sum_of_squares)
vmap_output = vectorized_sum_of_squares(batch_of_vectors)
print(f"vmap output: {vmap_output}") # Note: the print inside the function ran only once
Key points:
- Concise:
sum_of_squareswasn't rewritten; it was just wrapped injax.vmap; - Direct application: the whole batch of shape
(4, 3)is passed in directly; - Correct result:
[14. 77. 194. 1.]matches the manual loop; - Shape inference: the original function takes
(3,)and returns().vmapmaps over axis 0 of the input. Semantically that's like running the function once on each of the 4 length-3 vectors and stacking the 4 scalar results along a new axis 0, giving(4,). printran only once: becausevmapworks by tracing. It callssum_of_squaresonce with an abstract value that represents the whole batch (aBatchTracer). Each time it reaches a JAX primitive, it applies that primitive's "batching rule" on the spot and swaps in the batched version. Python'sprintis a side effect that runs at trace time, so it runs only once. The(3,)it prints is the single-example shape the function sees inside. To print values per example at run time, usejax.debug.print.
⚠️ Note: "
vmap"understood the batching pattern and optimized execution". It's simply a consequence of tracing. Pythonjax.debug.printand it prints once per example (see the code below). The function itself was not "run fewer times".
Default behavior: by default vmap maps over axis 0 of every input argument and stacks the results along axis 0 of the output. That is exactly the most common batching convention in machine learning.
Also, in JAX, calling a function repeatedly in a Python loop and then combining it with jit can lead to inefficient tracing (see Chapter 2: the loop gets unrolled, and tracing and compile time grow with the number of iterations).
Semantic equivalence (not how it's implemented):
jax.vmap(f)(xs) ≡ jnp.stack([f(x) for x in xs])
🔍 Extra: Python print vs. jax.debug.print¶
🟢 Must-read
- Python
print: runs once at trace time and sees aBatchTracerwith the single-example shape; jax.debug.print: runs at run time, and under vmap prints once per example;- Consequence: Python logic in the function body that depends on concrete values (
if,float(x), etc.) doesn't work under vmap (see the extra later in this section).
def show(v):
print("Python print (runs at trace time):", type(v).__name__, v.shape)
jax.debug.print("jax.debug.print (runs at run time): {}", v)
return jnp.sum(v ** 2)
import io, contextlib
buf = io.StringIO()
with contextlib.redirect_stdout(buf): # Capture the output so we can check how many times things were printed
res = jax.vmap(show)(batch_of_vectors[:3]).block_until_ready()
jax.effects_barrier() # Wait for all debug.print callbacks to finish
captured = buf.getvalue()
print(captured, "Result:", res)
# Check: Python print runs once at trace time; jax.debug.print runs once for each of the 3 examples
assert captured.count("Python print") == 1
assert captured.count("jax.debug.print") == 3
print("✔ Assertion passed: printing only once is caused by tracing; no computation was \"optimized away\"")
🔍 Extra: use jax.make_jaxpr to see exactly what vmap does¶
⚪ Optional: read this if you want to see what the program looks like after the vmap transformation.
jax.make_jaxpr prints the intermediate representation (jaxpr) that JAX gets from tracing. Comparing the jaxprs of f and vmap(f) shows that vmap(f) contains no loop at all. The arrays just have one more dimension, and each operation is swapped for its batched version (for example, reduce_sum now sums along axis 1, and the dot product dot_general gains a batch dimension). This shows directly that vmap transforms the program itself; it doesn't write a for loop for you.
def sos(v): # sum_of_squares without the print
return jnp.sum(v ** 2)
print("===== jaxpr of f =====")
print(jax.make_jaxpr(sos)(jnp.ones(3)))
print("===== jaxpr of vmap(f) =====")
jp_v = jaxpr_str(jax.make_jaxpr(jax.vmap(sos))(batch_of_vectors))
print(jp_v)
# Check: the vmapped program contains no loop primitives; the shapes just gain one dimension
assert not any(p in jp_v for p in ("while", "scan", "fori")) and "f32[4]" in jp_v
def dot(x, w):
return jnp.dot(x, w)
print("===== jaxpr of vmap(dot): dot_general gains batch dimensions ([0], [0]) =====")
print(jax.make_jaxpr(jax.vmap(dot))(jnp.ones((5, 3)), jnp.ones((5, 3))))
🔍 Extra: no Python if on batched data inside vmap¶
🟢 Must-read
Under vmap, the values inside the function are tracers that stand for "a whole batch of values". A Python if x > 0: can't decide which branch to take and raises TracerBoolConversionError (essentially the same restriction as under jit). Use jnp.where or lax.cond instead (4.7 covers what lax.cond costs under batching).
def relu_py(x):
if x > 0: # Python control flow depends on the mapped value
return x
return 0.0
try:
jax.vmap(relu_py)(jnp.array([-1.0, 2.0, 3.0]))
raised = False
except jax.errors.TracerBoolConversionError as e:
raised = True
print(type(e).__name__, ":", str(e).splitlines()[0])
assert raised, "A Python if that depends on batched data should raise"
# The right way: express the elementwise choice with jnp.where
relu = lambda x: jnp.where(x > 0, x, 0.0)
out = jax.vmap(relu)(jnp.array([-1.0, 2.0, 3.0]))
assert jnp.array_equal(out, jnp.array([0.0, 2.0, 3.0]))
print(out)
4.3 Mapping over specific arguments (in_axes, out_axes)¶
By default vmap treats axis 0 of every argument as the batch dimension. In practice you often need finer control: maybe only some arguments are batched (typically, model parameters are shared and only the data is batched), or the batch dimension isn't on axis 0. in_axes and out_axes give you precise control over how inputs and outputs are handled.
in_axes: which axis to map over for each input
- As a tuple, its length must equal the number of positional arguments passed at call time, with one element per positional argument. Keyword arguments are not covered by the tuple and are always mapped over axis 0;
- As a single integer (or
None), it applies to all positional arguments (the default isin_axes=0); - An integer
i: axisiof this argument is the batch dimension, andvmapprocesses it slice by slice along that axis; None: this argument is not mapped; it is broadcast whole and reused as is in every call;- All mapped arguments must have the same size along their mapped axes, or you get an error;
- Even if you map over axis 1 of an input, the batch dimension of the output is still on axis 0 by default.
out_axes: where the batch dimension goes in the output
- By default the results are stacked along axis 0 of the output;
out_axes=1puts the batch dimension on axis 1 of the output;- When the function returns several values (a pytree, such as a dict),
out_axescan be a pytree of the same structure, giving a batch axis for each returned value.
Shape sketch: in f(W, x), W has shape $(m, n)$ and x has shape $(n,)$:
$$ \text{vmap}(f,\ \text{in\_axes}=(\text{None}, 0))\big(W_{(m,n)},\ X_{(B,n)}\big) \to Y_{(B,m)}, \quad Y_b = W x_b $$
With in_axes and out_axes together, a single-example function can adapt to all kinds of batch layouts without changing its core logic or reshuffling dimensions by hand.
⚠️ Note: the length of an
in_axestuple must match the number of positional arguments at call time, not loosely "the number of function parameters". Keyword arguments are always mapped over axis 0 and the tuple doesn't cover them (the error is demonstrated at the end of this section). A list also works in 0.12, but the official docs say tuple, so stick with tuples.
Course code¶
def add_scalar(vector, scalar):
return vector + scalar
vectors = jnp.arange(12).reshape(4, 3) # 4 vectors, each of length 3
scalar_val = 100.0
# Map over axis 0 of vectors; broadcast scalar (not mapped)
vectorized_add_scalar = jax.vmap(add_scalar, in_axes=(0, None))
result = vectorized_add_scalar(vectors, scalar_val)
print("in_axes=(0, None):", result.shape)
print(result)
# Add a different scalar to each vector: map both arguments over axis 0
scalars = jnp.array([100., 200., 300., 400.])
result_batched_scalars = jax.vmap(add_scalar, in_axes=(0, 0))(vectors, scalars)
print("in_axes=(0, 0):", result_batched_scalars.shape)
print(result_batched_scalars)
# Map over a non-zero axis: after transposing, each column is a vector
vectors_T = vectors.T # shape (3, 4)
result_axis1 = jax.vmap(add_scalar, in_axes=(1, None))(vectors_T, scalar_val)
print("in_axes=(1, None):", result_axis1.shape, "(the output's batch dimension is still axis 0)")
print(result_axis1)
assert result_axis1.shape == (4, 3) and jnp.array_equal(result_axis1, result)
def process_vector(v):
return v * 2
# Default out_axes=0
out0 = jax.vmap(process_vector, in_axes=0, out_axes=0)(vectors)
print("out_axes=0:", out0.shape)
print(out0)
# out_axes=1: the original vector dimension (size 3) becomes output axis 0, and the batch dimension (size 4) goes to axis 1
out1 = jax.vmap(process_vector, in_axes=0, out_axes=1)(vectors)
print("out_axes=1:", out1.shape)
print(out1)
# Several return values (a pytree): a different out_axes for each output
def process_vector_pytree(v):
return {'sum': v.sum(), 'doubled': v * 2}
res = jax.vmap(process_vector_pytree, in_axes=0, out_axes={'sum': 0, 'doubled': 1})(vectors)
print("sum:", res['sum'].shape, res['sum'])
print("doubled:", res['doubled'].shape)
print(res['doubled'])
🔍 Extra: a more typical case, shared weights with per-example data¶
⚪ Optional: read this if you want to see a shared weight matrix and negative axes.
The examples above share a scalar. In practice it's more common to share model parameters. In the linear layer W @ x below, W is not mapped and x is mapped per example. Negative axes work too (-1 means the last axis).
k1, k2 = jax.random.split(jax.random.key(1))
W = jax.random.normal(k1, (3, 4)) # Shared weights
X = jax.random.normal(k2, (5, 4)) # 5 examples
def linear(W, x): # Single example: matrix times vector
return W @ x
Y = jax.vmap(linear, in_axes=(None, 0))(W, X)
print("Y shape:", Y.shape, "| matches the hand-written batched version:", jnp.allclose(Y, X @ W.T, atol=1e-5))
print("in_axes=-1 (data stored column-wise):", jnp.allclose(jax.vmap(linear, in_axes=(None, -1))(W, X.T), Y))
🔍 Extra: in_axes can be a pytree too (e.g. a dict)¶
⚪ Optional: read this if an argument is a dict and only some of its fields vary per example.
We showed a pytree out_axes above, and in_axes supports pytrees as well: each element can be a pytree with the same structure as the corresponding argument, or a "prefix" of it. When an argument is a dict, you can have some fields vary per example and share the others.
Rules:
- If an argument's
in_axesis an integer orNone, it applies to every leaf of that argument's whole pytree (the "prefix" rule); - If it's a dict, its keys must match the argument dict's keys exactly, or you get an error.
params = {"w": jnp.array([1.0, 2.0, 3.0]), # Shared weights: shape (3,)
"b": jnp.array([0.0, 10.0, 20.0, 30.0])} # One bias per example: shape (4,)
xb = jnp.arange(12.0).reshape(4, 3) # 4 examples
def predict_d(p, x):
return jnp.dot(p["w"], x) + p["b"]
# Use a dict in_axes for params: w is not mapped, b is mapped over axis 0
out = jax.vmap(predict_d, in_axes=({"w": None, "b": 0}, 0))(params, xb)
print(out)
print("Checked by hand:", xb @ params["w"] + params["b"])
# A structure mismatch (wrong key name) raises
try:
jax.vmap(predict_d, in_axes=({"w": None, "bias": 0}, 0))(params, xb)
except Exception as e:
print(type(e).__name__, ":", str(e)[:420])
🔍 Extra: the keyword-argument trap, out_axes=None, axis_size, and sum_match (new in 0.12)¶
⚪ Optional: read this if you need to pass keyword arguments to a vmapped function, want an output without a batch dimension, or the batch size can't be inferred.
- Keyword arguments are always mapped over axis 0: a tuple
in_axesdoesn't cover them. If you pass an argument that should be shared as a keyword argument, it gets mapped too, which may raise an error or produce a bug that's hard to spot. Recommendation: pass only positional arguments to a vmapped function. out_axes=None: means this output has no batch dimension. It's only valid if the output really doesn't vary across examples; otherwise it raises.axis_size: if every input hasin_axes=None(or there are no inputs at all),vmapcan't infer the batch size, and you have to give it explicitly withaxis_size.sum_match(a new parameter in the JAX 0.12vmapsignature, absent from older tutorials): when set toTrue, anout_axes=Noneoutput that varies across examples is summed along the batch dimension instead of raising. This is useful in autodiff, because "summing along the mapped axis" is exactly the transpose of "broadcasting along the mapped axis".
add = lambda a, b: a + b
# 1) Keyword arguments: a tuple in_axes doesn't cover them, so it raises with a hint
try:
jax.vmap(add, in_axes=(0, None))(jnp.ones(3), b=jnp.ones(3))
raised = False
except ValueError as e:
raised = "keyword arguments" in str(e)
print("Keyword argument error:", str(e)[:260], "...")
assert raised, "A tuple in_axes doesn't cover keyword arguments, so this should raise"
# Keyword arguments get mapped over axis 0: passing a length-3 b gives elementwise addition, not "broadcast as a whole"
assert jnp.array_equal(jax.vmap(add)(jnp.arange(3.0), b=jnp.arange(3.0)), 2 * jnp.arange(3.0))
# 2) out_axes=None but the output depends on batched data -> error
try:
jax.vmap(lambda x: 2 * x, out_axes=None)(jnp.arange(3.0))
except ValueError as e:
print("out_axes=None error:", e)
# When the output doesn't depend on batched data, out_axes=None is valid
print("Valid use of out_axes=None:", jax.vmap(lambda x, c: (x * c, c), in_axes=(0, None), out_axes=(0, None))(jnp.arange(3.0), 5.0))
# 3) axis_size: required when no input is mapped
try:
jax.vmap(lambda c: c * jnp.ones(2), in_axes=None)(3.0)
except ValueError as e:
print("Missing axis_size error:", e)
print("axis_size=4:\n", jax.vmap(lambda c: c * jnp.ones(2), in_axes=None, axis_size=4)(3.0))
# 4) sum_match=True: an out_axes=None output is summed along the batch dimension
sm = jax.vmap(lambda x: 2.0 * x, out_axes=None, sum_match=True)(jnp.arange(3.0))
assert sm == 6.0 # 2*(0+1+2)
print("sum_match:", sm)
4.4 Handling multiple batched arguments¶
Functions you want to vectorize often have several inputs: for example, pairwise dot products between two batches of vectors, or multiplying a batch of vectors by the same matrix.
Default behavior: in_axes defaults to 0, which means all positional arguments are mapped over axis 0. So as soon as one argument shouldn't be mapped (a scalar, say), you have to write out the tuple explicitly, e.g. in_axes=(0, None).
⚠️ Note: the default
in_axes=0maps every positional argument, not just the first one. Applyingjax.vmap(scale_vector)directly toscale_vector(batch, 2.0)raisesValueErrorbecause the scalar has no axis 0 (checked withassertbelow).
Each element of in_axes decides how the corresponding argument takes part in the mapping:
- An integer
i: map over axisiof this argument; the size of that axis is the batch size. When several arguments use integerin_axes, the sizes of the axes they name must match; None: not mapped; it's broadcast as a constant into every call.
vmap takes the $i$-th slice from every mapped argument at the same time and hands them to f together:
$$ \text{vmap}(f,\ \text{in\_axes}=(0, \text{None}, 0))(X, y, Z)_i = f(X_i,\ y,\ Z_i) $$
The course illustrates this with a diagram: vmap takes $x_0, x_1, \dots$ from xs (along axis 0) and $z_0, z_1, \dots$ from zs (along axis 0), passes the same y_val to every call, and finally stacks all the results along axis 0 of the output into a batch. Note that this is pairing by position, not all combinations (for all combinations, use the nested vmap in 4.5).
There are three common scenarios:
- Map some arguments, broadcast the rest: e.g.
in_axes=(0, None); - Map several arguments along the same axis: e.g.
in_axes=(0, 0); the batch size comes from axis 0 of each argument, and they must be equal; - Map different arguments along different axes (advanced): e.g.
in_axes=(1, 0). Note: whichever axis an input is mapped over, the batch dimension goes on axis 0 of the output by default. To put it elsewhere, useout_axes.
⚠️ Note: in the "column scaling" example of scenario 3, it's easy to assume that
in_axes=(1, 0)puts the mapped dimension back on axis 1, giving a 3×4 layout. In fact the default isout_axes=0, so the result has shape(4, 3), which is exactly the transpose of the 3×4 layout. To get 3×4 you must addout_axes=1(see the ❌/✅ comparison below).
Course code¶
# Scenario 1: map some arguments, broadcast the others
def scale_vector(vector, scalar):
return vector * scalar
vector = jnp.arange(3.)
scalar = 2.0
print("Single scaling:", scale_vector(vector, scalar)) # [0. 2. 4.]
batch_of_vectors = jnp.arange(9.).reshape((3, 3))
single_scalar = 2.0
vectorized_scale = jax.vmap(scale_vector, in_axes=(0, None))
scaled = vectorized_scale(batch_of_vectors, single_scalar)
print("Batched scaling:\n", scaled)
assert jnp.array_equal(scaled, batch_of_vectors * 2.0)
# ❌ Wrong: assuming only the first argument is mapped by default, and leaving out in_axes
try:
jax.vmap(scale_vector)(batch_of_vectors, 2.0)
raised = False
except ValueError as e:
raised = True
print("Measured ValueError:", e)
assert raised, "The default in_axes=0 also tries to map the scalar argument, so this should raise"
# ✅ Recommended: the default in_axes=0 maps all positional arguments, so a shared scalar needs an explicit None
ok = jax.vmap(scale_vector, in_axes=(0, None))(batch_of_vectors, 2.0)
assert jnp.array_equal(ok, batch_of_vectors * 2.0)
# Check "the default maps all positional arguments": when both arguments have a batch dimension, the default in_axes maps both
def add_vectors(vec1, vec2):
return vec1 + vec2
b1 = jnp.arange(9.).reshape(3, 3)
b2 = jnp.arange(3.)[:, None] * jnp.ones((3, 3)) # row i is all i
default_out = jax.vmap(add_vectors)(b1, b2)
assert jnp.array_equal(default_out, b1 + b2) # example i = b1[i] + b2[i]; the second argument was mapped too
print("✔ Assertion passed: the default in_axes=0 maps all positional arguments")
# Scenario 2: map several arguments along the same axis
batch1 = jnp.arange(9.).reshape((3, 3))
batch2 = jnp.ones((3, 3)) * 10
vectorized_add = jax.vmap(add_vectors, in_axes=(0, 0))
added = vectorized_add(batch1, batch2)
print(added)
assert jnp.array_equal(added, batch1 + 10)
# Scenario 3 (first half): row × filter vector; map over the matrix rows, broadcast the filter vector
def apply_filter(row_vector, filter_vector):
return row_vector * filter_vector
matrix = jnp.arange(12.).reshape((3, 4))
filter_v = jnp.array([1., 10., 100., 1000.])
filtered = jax.vmap(apply_filter, in_axes=(0, None))(matrix, filter_v)
print("Row filtering in_axes=(0, None):\n", filtered)
assert filtered.shape == (3, 4)
# ❌ Easy to misread: expecting a 3x4 result with "the mapped dimension put back on axis 1"
def scale_column(column_vector, scalar):
return column_vector * scalar
scalars = jnp.array([1., 10., 100., 1000.])
if matrix.shape[1] == scalars.shape[0]:
scaled_cols = jax.vmap(scale_column, in_axes=(1, 0))(matrix, scalars)
print("Measured shape:", scaled_cols.shape)
print(scaled_cols)
# Check: default out_axes=0, so the batch dimension (4 columns) is on output axis 0 -> (4, 3), the transpose of what we wanted
assert scaled_cols.shape == (4, 3)
assert jnp.allclose(scaled_cols, (matrix * scalars).T)
# ✅ Recommended: use out_axes=1 to put the batch dimension back on axis 1 and get the 3x4 layout you want
scaled_cols_fixed = jax.vmap(scale_column, in_axes=(1, 0), out_axes=1)(matrix, scalars)
print("Fixed shape:", scaled_cols_fixed.shape)
print(scaled_cols_fixed)
assert scaled_cols_fixed.shape == (3, 4)
assert jnp.allclose(scaled_cols_fixed, matrix * scalars) # matches the broadcasting version
print("✔ Assertion passed")
🔍 Extra: the two most common errors and how to track them down¶
⚪ Optional: look here when you hit a shape-related vmap error.
- Mismatched batch sizes: two mapped arguments have different lengths along their mapped axes. Note that
vmapdoes not stretch length-1 axes the way NumPy broadcasting does. For a shared argument, useNoneinstead of adding a length-1 axis to it. in_axesdoesn't match the argument structure: the tuple length doesn't match the number of positional arguments, or the mapped axis is beyond the array's number of dimensions.
The error messages in JAX 0.12 are quite specific (they tell you which argument, which axis and what size), so read the error carefully first.
f2 = lambda a, b: a + b
cases = {
"Mismatched batch sizes (3 vs 4)": lambda: jax.vmap(f2)(jnp.ones(3), jnp.ones(4)),
"Length 1 is not broadcast automatically": lambda: jax.vmap(f2)(jnp.ones((3, 2)), jnp.ones((1, 2))),
"in_axes length doesn't match the number of arguments": lambda: jax.vmap(f2, in_axes=(0,))(jnp.ones(3), jnp.ones(3)),
"Mapped axis beyond the array's dimensions": lambda: jax.vmap(lambda a: a, in_axes=1)(jnp.ones(3)),
}
n_raised = 0
for name, fn in cases.items():
try:
fn()
print(f"[{name}] no error?!")
except ValueError as e:
n_raised += 1
print(f"[{name}]\n ValueError: {e}\n")
assert n_raised == len(cases)
4.5 Nested vmap¶
Data can have several levels of batch structure. For example, a batch of sentences where each sentence is a sequence of words, and you want to operate on each word vector; or you want the relationship between "every element of one batch" and "every element of another batch". vmap is itself a function transformation, so you can wrap it in another layer and map over several dimensions at once.
Why nest? A single vmap(f, in_axes=(0, 0)) can only pair two batches by position. If you want all combinations of "every matrix × every vector" (a bit like an "outer product" at the batch level), you need to map the matrix batch and the vector batch independently, and that takes two layers of vmap.
How to read a nested vmap (from the inside out):
- The inner
vmap(f, in_axes=(ax1, ax2))gives a new function that mapsfover axesax1andax2of its own inputs; - The outer
vmap(..., in_axes=(ax3, ax4))then maps that new function over axesax3andax4of the outer inputs; - An argument with
Noneat some layer is not mapped at that layer and is broadcast as is.
vmap(vmap(f, in_axes=(None, 0)), in_axes=(0, None)) is the classic pattern for computing pairwise relationships between two batches:
$$ M_{ij} = f(x_i, y_j),\qquad X \in \mathbb{R}^{M\times d},\ Y \in \mathbb{R}^{N\times d} \;\Rightarrow\; M \in \mathbb{R}^{M \times N} $$
In the output, the outer batch dimension comes first and the inner one second. Nested vmap also composes with jit and grad. Remembering that "each vmap adds one layer of mapping" is the key to writing complex batching transformations.
Course code¶
# One matrix times one vector
def matvec_multiply(matrix, vector):
return jnp.dot(matrix, vector)
matrix = jnp.arange(6.).reshape(2, 3)
vector = jnp.arange(3.)
print("Single matvec:", matvec_multiply(matrix, vector)) # [ 5. 14.], shape (2,)
# Single vmap: matrix batch and vector batch paired by position
batch_of_matrices = jnp.arange(18.).reshape(3, 2, 3) # 3 matrices of size 2x3
batch_of_vectors = jnp.arange(9.).reshape(3, 3) # 3 vectors of dimension 3
batched_matvec = jax.vmap(matvec_multiply, in_axes=(0, 0))
print("Paired matvec shape:", batched_matvec(batch_of_matrices, batch_of_vectors).shape) # (3, 2)
# Nested vmap: every matrix × every vector (all combinations)
batch_of_vectors_outer = jnp.arange(12.).reshape(4, 3) # 4 vectors
# Inner in_axes=(None, 0): matrix fixed, loop over all vectors
# Outer in_axes=(0, None): loop over all matrices, share the whole vector batch
pairwise_matvec = jax.vmap(
jax.vmap(matvec_multiply, in_axes=(None, 0)),
in_axes=(0, None),
)
result = pairwise_matvec(batch_of_matrices, batch_of_vectors_outer)
print("Pairwise matvec shape:", result.shape) # (3, 4, 2)
assert result.shape == (3, 4, 2)
# Check: result[i, j] should equal matrix i times vector j
assert jnp.allclose(result[1, 2], matvec_multiply(batch_of_matrices[1], batch_of_vectors_outer[2]))
print("✔ result[1, 2] == matrix 1 @ vector 2")
# Pairwise distances (squared Euclidean distance)
def euclidean_distance_sq(point_a, point_b):
return jnp.sum((point_a - point_b)**2)
points_a = jnp.array([[1., 0.], [0., 1.], [-1., 0.]]) # (3, 2)
points_b = jnp.array([[2., 2.], [-2., 2.], [2., -2.], [-2., -2.]]) # (4, 2)
inner_map = jax.vmap(euclidean_distance_sq, in_axes=(None, 0)) # one a against all b
pairwise_distance_sq = jax.vmap(inner_map, in_axes=(0, None)) # all a
# Equivalent one-liner
pairwise_distance_sq_direct = jax.vmap(
jax.vmap(euclidean_distance_sq, in_axes=(None, 0)), in_axes=(0, None))
distance_matrix = pairwise_distance_sq(points_a, points_b)
print("Distance matrix shape:", distance_matrix.shape) # (3, 4)
print(distance_matrix)
assert jnp.allclose(distance_matrix, pairwise_distance_sq_direct(points_a, points_b))
assert jnp.array_equal(distance_matrix, jnp.array([[5., 13., 5., 13.], [5., 5., 13., 13.], [13., 5., 13., 5.]]))
print("Check [0, 0] by hand:", euclidean_distance_sq(points_a[0], points_b[0])) # (1-2)^2 + (0-2)^2 = 5
# Combining nested vmap with jit and grad
jitted_pairwise_distance_sq = jax.jit(pairwise_distance_sq_direct)
print("jit version matches:", jnp.allclose(jitted_pairwise_distance_sq(points_a, points_b), distance_matrix))
def sum_of_distances(p_a, p_b):
dist_matrix = pairwise_distance_sq_direct(p_a, p_b)
return jnp.sum(dist_matrix)
grad_sum_dist_wrt_a = jax.grad(sum_of_distances, argnums=0)
gradients_a = grad_sum_dist_wrt_a(points_a, points_b)
print("Gradient w.r.t. points_a, shape:", gradients_a.shape) # (3, 2), same as points_a
print(gradients_a)
# Check against the analytic result: d/da_i sum_j |a_i - b_j|^2 = sum_j 2 (a_i - b_j)
assert jnp.allclose(gradients_a, 2 * (points_b.shape[0] * points_a - points_b.sum(0)))
print("✔ The gradient matches the analytic result")
⚠️ Note: the "pairwise matvec" code is easy to read as "first map over the matrices (axis 0), then over the vectors (axis 0)". What it actually means is: the inner
(None, 0)loops over the vectors, the outer(0, None)loops over the matrices, and the output is(number of matrices, number of vectors, 2)=(3, 4, 2). The code comments above spell this out, andassertchecks both the shape and the elements.
🔍 Extra: nesting order decides the order of output axes; several batch dimensions¶
⚪ Optional: read this if you need several batch dimensions or care about the order of the output axes.
Swap what the inner and outer layers map over, and the output becomes the transpose. For data with two batch dimensions, like (B, T, d), two layers of vmap give a (B, T) result directly.
k1, k2 = jax.random.split(jax.random.key(3))
Xs = jax.random.normal(k1, (4, 3)) # M=4
Ys = jax.random.normal(k2, (5, 3)) # N=5
G = jax.vmap(jax.vmap(jnp.dot, in_axes=(None, 0)), in_axes=(0, None))(Xs, Ys)
G2 = jax.vmap(jax.vmap(jnp.dot, in_axes=(0, None)), in_axes=(None, 0))(Xs, Ys)
print("G:", G.shape, "| G2 with the order swapped:", G2.shape, "| G2 == G.T:", jnp.allclose(G2, G.T))
print("equals X @ Y.T:", jnp.allclose(G, Xs @ Ys.T, atol=1e-5))
data = jnp.ones((2, 3, 4)) # (B, T, d)
print("Two batch dimensions, norm at each position:", jax.vmap(jax.vmap(jnp.linalg.norm))(data).shape)
🔍 Extra: numerical stability of distance matrices, and the gradient trap of sqrt at 0¶
⚪ Optional: read this if you need gradients of a distance matrix or write the expanded formula by hand.
So far we used "squared distance", so the gradient is fine. But the exercises in 4.8 use the true Euclidean distance $\sqrt{\cdot}$, and then two things need care:
- Numerical stability: hand-vectorized code often uses the expansion $\lVert x\rVert^2 + \lVert y\rVert^2 - 2x^\top y$. Because of floating-point cancellation it can produce small negative numbers, which turn into
nanafter the square root, so you needclip. The vmap version computes $x - y$ directly and doesn't have this problem (at the cost of creating an $(M, N, d)$ intermediate array; see the memory discussion in 4.7). - The gradient is
nanwhere the distance is 0: $\frac{\partial}{\partial x}\sqrt{s} = \frac{x-y}{\sqrt{s}}$, which is $0/0$ when $x=y$, for example when you differentiate the distance matrix of a set of points with itself (zeros on the diagonal). The fix is the "doublewhere"; for the idea and the general pattern see Chapter 3, extra in 3.8 "NaN gradients and double where", and for more advanced material see 07 topic 2.
def dist(x, y):
return jnp.sqrt(jnp.sum((x - y) ** 2))
pdist = jax.vmap(jax.vmap(dist, in_axes=(None, 0)), in_axes=(0, None))
D = pdist(Xs, Ys)
# For comparison: the broadcasting version (expanded formula) needs clip to avoid negative numbers
sq = jnp.sum(Xs**2, 1)[:, None] + jnp.sum(Ys**2, 1)[None, :] - 2 * Xs @ Ys.T
D_ref = jnp.sqrt(jnp.clip(sq, 0.0))
assert jnp.allclose(D, D_ref, atol=1e-4)
print("Distance matrix shape:", D.shape, "| matches the expanded formula")
# Gradient trap: the gradient is nan when x == y (the "double where" fix is in Chapter 3, 3.8)
g_naive = jax.grad(dist)(jnp.ones(3), jnp.ones(3))
print("Gradient of naive dist at x==y:", g_naive)
assert jnp.all(jnp.isnan(g_naive))
🔍 Extra: batched matrix multiplication with vmap, compared with jnp.einsum and @¶
⚪ Optional: read this when choosing between vmap, einsum and @.
To compute $A_b B_b$ for a batch of matrices $A_b \in \mathbb{R}^{n\times k}$, $B_b \in \mathbb{R}^{k\times m}$, there are three ways to write it:
jax.vmap(jnp.matmul)(A, B): write a single matrix product and leave batching to vmap;jnp.einsum('bnk,bkm->bnm', A, B): spell out the dimension relationships with subscripts;A @ B:matmulalready treats leading dimensions as batch dimensions.
At the jaxpr level all three are the same dot_general with a batch dimension (below, assert compares the three jaxprs character by character), so their speed is essentially the same; timing differences of a fraction of a millisecond are just measurement noise. Pick by readability: with many dimensions, transposes and contractions, einsum is clearer; when the function is complex (more than a single matrix product), vmap shines.
k1, k2 = jax.random.split(jax.random.key(4))
Ab = jax.random.normal(k1, (64, 128, 96))
Bb = jax.random.normal(k2, (64, 96, 80))
bmm_vmap = jax.jit(jax.vmap(jnp.matmul))
bmm_einsum = jax.jit(lambda a, b: jnp.einsum("bnk,bkm->bnm", a, b))
bmm_at = jax.jit(lambda a, b: a @ b)
r1, r2, r3 = bmm_vmap(Ab, Bb), bmm_einsum(Ab, Bb), bmm_at(Ab, Bb)
assert jnp.allclose(r1, r2, atol=1e-3) and jnp.allclose(r1, r3, atol=1e-3)
print("Result shape:", r1.shape, "| all three agree")
for name, fn in [("vmap(matmul)", bmm_vmap), ("einsum", bmm_einsum), ("a @ b", bmm_at)]:
print(f"{name:14s} {bench(fn, Ab, Bb):.3f} ms")
# Look at the jaxpr: all are dot_general with a batch dimension
jp_bmm = [jaxpr_str(jax.make_jaxpr(f)(Ab[:2], Bb[:2])) for f in
(jax.vmap(jnp.matmul), lambda a, b: jnp.einsum("bnk,bkm->bnm", a, b), lambda a, b: a @ b)]
assert jp_bmm[0] == jp_bmm[1] == jp_bmm[2] # Check: the jaxprs of all three versions are identical
print(jax.make_jaxpr(jax.vmap(jnp.matmul))(Ab[:2], Bb[:2]))
print(jax.make_jaxpr(lambda a, b: jnp.einsum("bnk,bkm->bnm", a, b))(Ab[:2], Bb[:2]))
4.6 Combining vmap with jit and grad¶
One of JAX's big strengths is the composability of its function transformations: jit (compilation), grad (differentiation) and vmap (vectorization) all take a function and return a function, so you can combine them freely. Put together, they let you take gradients over a batch of data and compile the whole computation into code that runs efficiently on an accelerator.
1. vmap + jit. The recommended form is jit(vmap(f)). vmap(f) first rewrites the primitives in f as batched operations, and jit then uses XLA to compile that batched function as a whole into kernels optimized for CPU/GPU/TPU. vmap(jit(f)) works too: inside JAX, jit is itself a primitive, and vmap pushes the batch dimension inside the jit, so the compiled kernels are essentially the same as for jit(vmap(f)). The only difference is on the Python side: the outermost layer of vmap(jit(f)) is an "eagerly executed vmap", which redoes the batching transformation in Python on every call, while jit(vmap(f)) traces only on the first call. So putting jit on the outside has the least overhead.
⚠️ Note: it's easy to believe that
vmap(jit(f))"compiles the inner f first and then vectorizes the calls to the compiled function, so the vectorization logic lives outside the compiled kernel and may be less efficient". In fact the jaxpr shows the batch dimension pushed inside thejit(thejit[...]already operates on the whole batch). For large batches both orders take the same time; for small batchesvmap(jit(f))is slightly slower because of Python-side overhead, not because "the vectorization is outside the kernel" (jaxpr and timings below).
2. vmap + grad. In machine learning you often need gradients for a batch of examples rather than a single one. The approach: get a per-example gradient function with grad, then vectorize it with vmap: vmap(grad(loss), in_axes=(None, 0, 0)) (parameters shared, data mapped per example). The per-example gradients carry the batch dimension that vmap adds.
The more common need is the batch-averaged gradient. There are two ways: compute per-example gradients and then average, or define the batch-averaged loss and then take its gradient. The two are mathematically equivalent:
$$ \nabla_\theta \frac{1}{N}\sum_i \ell(\theta; x_i, y_i) = \frac{1}{N}\sum_i \nabla_\theta \ell(\theta; x_i, y_i) $$
How to choose: when you really need per-example results, use vmap(grad(f), ...), which states the intent more clearly. When you only need the batch-averaged gradient, use grad(mean_loss_fn). That's the standard approach for gradient descent, and it doesn't create a "batch size × number of parameters" intermediate. Note that grad requires a scalar output, so you can't write grad(vmap(loss)) directly; sum or average first.
3. vmap + grad + jit. This is the basic pattern for training models efficiently in JAX, in four steps: define the loss → take the gradient with grad → vectorize with vmap → compile with jit.
- Batch-averaged gradient:
jit(grad(mean_loss_fn)); - Per-example gradients:
jit(vmap(grad(single_loss_fn))).
Things to keep in mind:
- Debugging: once transformations are stacked, error tracebacks can go deep into JAX internals. Use
jax.debug.print, or temporarily removejitto locate the problem; - Memory:
vmapover very large batches can use a lot of memory; be especially careful with batch sizes on GPUs/TPUs with limited memory (the memory math for per-example gradients is in this section's special topic); - Order: the order in which you apply transformations affects both correctness and performance.
Random numbers: the examples in this section generate random data. JAX 0.12 recommends creating typed keys with jax.random.key(seed) and using each key only once. When you need several sets of random numbers, jax.random.split first.
⚠️ Note: older code often creates a key with
jax.random.PRNGKey(42)and then uses the same key to generate both theuniformx and thenormalnoise.PRNGKeyis the old-style API; more importantly, the two sets of random numbers you get are highly correlated (linear correlation about 0.98, rank correlation 1), so the "noise" isn't independent at all (the ❌/✅ comparison below checks this withassert).
Course code¶
import time
# A simple linear model
def predict(params, x):
w, b = params
return w * x + b
params = (jnp.array(2.0), jnp.array(1.0)) # (w, b)
batch_x = jnp.arange(10000.0)
# Map over axis 0 of batch_x; params stays fixed
batched_predict_vmap = jax.vmap(predict, in_axes=(None, 0))
jitted_batched_predict = jax.jit(batched_predict_vmap)
t0 = time.perf_counter()
result_vmap = batched_predict_vmap(params, batch_x).block_until_ready()
print(f"vmap only, run time: {(time.perf_counter() - t0) * 1e3:.3f} ms")
t0 = time.perf_counter()
result_jit_vmap = jitted_batched_predict(params, batch_x).block_until_ready()
print(f"jit(vmap) run time (first run, includes compilation): {(time.perf_counter() - t0) * 1e3:.3f} ms")
t0 = time.perf_counter()
result_jit_vmap = jitted_batched_predict(params, batch_x).block_until_ready()
print(f"jit(vmap) run time (second run): {(time.perf_counter() - t0) * 1e3:.3f} ms")
assert jnp.allclose(result_vmap, result_jit_vmap)
print("Results match")
# The single timings above only show that "the first run includes compilation"; to compare steady-state times, warm up and average over several runs
print(f"Average after warm-up: vmap {bench(batched_predict_vmap, params, batch_x):.3f} ms | "
f"jit(vmap) {bench(jitted_batched_predict, params, batch_x):.3f} ms")
# grad with vmap: per-example gradients
def squared_error(params, x, y_true):
w, b = params
y_pred = w * x + b
loss = (y_pred - y_true)**2
return loss
# Gradient function for a single (x, y_true): returns (dL/dw, dL/db)
grad_loss_single = jax.grad(squared_error, argnums=0)
params = (2.0, 1.0)
batch_x = jnp.linspace(0, 1, 5)
batch_y_true = 2.5 * batch_x + 0.5 # True relationship: w=2.5, b=0.5
grad_loss_batch = jax.vmap(grad_loss_single, in_axes=(None, 0, 0))
batch_gradients = grad_loss_batch(params, batch_x, batch_y_true)
print("params:", params)
print("batch_x:", batch_x)
print("batch_y_true:", batch_y_true)
print("per-example dw:", batch_gradients[0], batch_gradients[0].shape) # (5,)
print("per-example db:", batch_gradients[1], batch_gradients[1].shape) # (5,)
# Check against the analytic result: dL/dw = 2 (w x + b - y) x, dL/db = 2 (w x + b - y)
resid = 2.0 * batch_x + 1.0 - batch_y_true
assert jnp.allclose(batch_gradients[0], 2 * resid * batch_x) and jnp.allclose(batch_gradients[1], 2 * resid)
# Batch-averaged loss -> averaged gradient
def mean_squared_error(params, batch_x, batch_y_true):
batch_predict = jax.vmap(predict, in_axes=(None, 0))
predictions = batch_predict(params, batch_x)
return jnp.mean((predictions - batch_y_true)**2)
grad_mean_loss = jax.grad(mean_squared_error, argnums=0)
mean_batch_gradient = grad_mean_loss(params, batch_x, batch_y_true)
print("Averaged gradient:", mean_batch_gradient)
print("Shapes:", jax.tree.map(jnp.shape, mean_batch_gradient)) # every component is a scalar ()
# Averaged gradient == average of the per-example gradients
assert jnp.allclose(mean_batch_gradient[0], jnp.mean(batch_gradients[0]))
assert jnp.allclose(mean_batch_gradient[1], jnp.mean(batch_gradients[1]))
print("✔ gradient of the mean loss == mean of the per-example gradients")
# ❌ Old style (outdated): PRNGKey is an old-style raw key, and the same key is used for two draws, so x and the "noise" are highly correlated
key = jax.random.PRNGKey(42)
print("PRNGKey measured:", key.dtype, key.shape, "(a raw uint32 array, not a typed key)")
large_batch_size = 100000
large_batch_x_bad = jax.random.uniform(key, (large_batch_size,))
noise_bad = jax.random.normal(key, (large_batch_size,)) # the same key again
rank = lambda a: np.argsort(np.argsort(np.asarray(a)))
corr = np.corrcoef(np.asarray(large_batch_x_bad), np.asarray(noise_bad))[0, 1]
rank_corr = np.corrcoef(rank(large_batch_x_bad), rank(noise_bad))[0, 1]
print(f"Measured: linear correlation between x and noise {corr:.4f}, rank correlation {rank_corr:.4f}")
assert corr > 0.95 and rank_corr > 0.999 # Check: uniform and normal draws from the same key are highly correlated
# ✅ Recommended: create a typed key with jax.random.key, split it, then use each part separately
key = jax.random.key(42)
key_x, key_noise = jax.random.split(key)
print("jax.random.key measured:", key.dtype, key.shape)
large_batch_x = jax.random.uniform(key_x, (large_batch_size,))
noise = jax.random.normal(key_noise, (large_batch_size,))
corr_ok = np.corrcoef(np.asarray(large_batch_x), np.asarray(noise))[0, 1]
print(f"Correlation after split: {corr_ok:.4f}")
assert abs(corr_ok) < 0.02 # Independent: correlation close to 0
large_batch_y_true = 2.5 * large_batch_x + 0.5 + 0.1 * noise
# Course code (continued): jit(vmap(grad(f))), per-example gradients for a large batch (data from the recommended version above)
vmap_grad_loss = jax.vmap(grad_loss_single, in_axes=(None, 0, 0))
jit_vmap_grad_loss = jax.jit(vmap_grad_loss)
params = (2.0, 1.0)
t0 = time.perf_counter()
per_example_grads = jit_vmap_grad_loss(params, large_batch_x, large_batch_y_true)
jax.block_until_ready(per_example_grads)
print(f"First run (includes compilation): {(time.perf_counter() - t0) * 1e3:.2f} ms")
t0 = time.perf_counter()
per_example_grads = jit_vmap_grad_loss(params, large_batch_x, large_batch_y_true)
jax.block_until_ready(per_example_grads)
print(f"Second run: {(time.perf_counter() - t0) * 1e3:.2f} ms")
print("dw shape:", per_example_grads[0].shape, "| db shape:", per_example_grads[1].shape)
assert per_example_grads[0].shape == per_example_grads[1].shape == (large_batch_size,)
🔍 Extra: checking that vmap(jit(f)) pushes the batch dimension inside the jit¶
⚪ Optional: read this if you want to verify the behavior of vmap(jit(f)) yourself with a jaxpr.
f_feat = lambda x: jnp.sin(x) @ jnp.cos(x)
jp = jaxpr_str(jax.make_jaxpr(jax.vmap(jax.jit(f_feat)))(jnp.ones((4, 3))))
print(jp)
# Check: inside the jit, the operations already act on the whole (4, 3) batch
inner = jp[jp.index("jaxpr={"):]
assert "f32[4,3] = sin" in inner and "f32[4] = dot_general" in inner
for shape in [(512, 256), (4096, 1024)]:
xb = jax.random.normal(jax.random.key(6), shape)
a_ = jax.jit(jax.vmap(f_feat))
b_ = jax.vmap(jax.jit(f_feat))
assert jnp.allclose(a_(xb), b_(xb), rtol=1e-4, atol=1e-3)
print(f"batch {shape}: jit(vmap(f)) {bench(a_, xb):.3f} ms | vmap(jit(f)) {bench(b_, xb):.3f} ms")
🔍 Special topic: per-example gradients¶
🟢 Must-read
How it works. Let the single-example loss be $\ell_i(\theta) = \ell(\theta; x_i, y_i)$. The per-example gradient is $g_i = \nabla_\theta \ell_i(\theta)$.
jax.grad(loss_single)builds a "forward + backward (VJP)" program for one example;jax.vmap(..., in_axes=(None, 0, 0))then batches that whole program: the parameters $\theta$ are shared (None) and the data is mapped per example. Since each example's gradient depends on its own $x_i$, every gradient leaf in the output gains a batch dimension, and its shape goes fromp.shapeto(B, *p.shape);- Another way to see it: stack the $B$ losses into a vector $L(\theta) \in \mathbb{R}^B$. The per-example gradients are then the Jacobian $J = \partial L / \partial \theta \in \mathbb{R}^{B \times P}$ (the same as what
jax.jacrevgives). A plaingrad(mean_loss)only computes $\frac{1}{B}\mathbf{1}^\top J$, a single VJP, where the batch dimension is summed away during backpropagation, so it uses far less memory.
In PyTorch, per-example gradients used to require dedicated libraries like Opacus or functorch (functorch was later merged into torch.func, where the idiom is also vmap(grad(...))). In JAX it has been a composition of two basic transformations from the start: one line, vmap(grad(loss)), is enough. Below we walk through the theory, a cross-check, timing, applications and memory with a two-layer MLP regression model.
# Setup: a two-layer MLP + regression data, with the labels of 5 examples deliberately corrupted (used later to find outliers)
def init_mlp_params(key, d_in, d_hidden):
k1, k2 = jax.random.split(key)
return {"W1": jax.random.normal(k1, (d_hidden, d_in)) / jnp.sqrt(d_in), "b1": jnp.zeros(d_hidden),
"w2": jax.random.normal(k2, (d_hidden,)) / jnp.sqrt(d_hidden), "b2": jnp.array(0.0)}
def mlp_loss(p, x, y): # Single-example loss (scalar)
h = jnp.tanh(p["W1"] @ x + p["b1"])
return (jnp.dot(p["w2"], h) + p["b2"] - y) ** 2
B, D_IN, H = 256, 8, 32
kp, kx, kn = jax.random.split(jax.random.key(100), 3)
pe_params = init_mlp_params(kp, D_IN, H)
X_pe = jax.random.normal(kx, (B, D_IN))
Y_pe = jnp.sin(X_pe[:, 0]) + 0.1 * jax.random.normal(kn, (B,))
bad_idx = jnp.array([3, 50, 99, 150, 222]) # "mislabeled" examples
Y_pe = Y_pe.at[bad_idx].add(20.0)
P = sum(l.size for l in jax.tree.leaves(pe_params))
print("Number of parameters P =", P, "| batch size B =", B)
# 1) Cross-check: vmap(grad) vs. a Python loop taking gradients one by one, and compare with jacrev (the Jacobian view)
per_ex_grad = jax.jit(jax.vmap(jax.grad(mlp_loss), in_axes=(None, 0, 0)))
single_grad = jax.jit(jax.grad(mlp_loss))
def loop_per_ex_grad(p, X, Y):
gs = [single_grad(p, X[i], Y[i]) for i in range(X.shape[0])]
return jax.tree.map(lambda *t: jnp.stack(t), *gs)
g_vmap = per_ex_grad(pe_params, X_pe, Y_pe)
g_loop = loop_per_ex_grad(pe_params, X_pe, Y_pe)
g_jac = jax.jacrev(lambda p: jax.vmap(mlp_loss, in_axes=(None, 0, 0))(p, X_pe, Y_pe))(pe_params)
tree_allclose = lambda a, b, tol=1e-5: all(jax.tree.leaves(jax.tree.map(lambda u, v: bool(jnp.allclose(u, v, atol=tol)), a, b)))
assert tree_allclose(g_vmap, g_loop)
assert tree_allclose(g_vmap, g_jac)
print("Per-example gradient shapes:", jax.tree.map(jnp.shape, g_vmap))
print("✔ vmap(grad) == looping over examples == jacrev")
t_loop = bench(lambda p, X, Y: loop_per_ex_grad(p, X, Y)["W1"], pe_params, X_pe, Y_pe, n=2)
t_vmap = bench(lambda p, X, Y: per_ex_grad(p, X, Y)["W1"], pe_params, X_pe, Y_pe)
print(f"Python loop: {t_loop:.2f} ms | jit(vmap(grad)): {t_vmap:.3f} ms | about {t_loop / t_vmap:.0f}x faster")
# 2) Mean of the per-example gradients == gradient of the mean loss
def mean_mlp_loss(p, X, Y):
return jnp.mean(jax.vmap(mlp_loss, in_axes=(None, 0, 0))(p, X, Y))
g_mean_loss = jax.jit(jax.grad(mean_mlp_loss))(pe_params, X_pe, Y_pe)
g_mean_of_per_ex = jax.tree.map(lambda t: t.mean(axis=0), g_vmap)
assert tree_allclose(g_mean_loss, g_mean_of_per_ex)
print("✔ mean_i g_i == grad(mean loss)")
# For contrast: grad(vmap(loss)) doesn't work directly, because the output isn't a scalar
try:
jax.grad(jax.vmap(mlp_loss, in_axes=(None, 0, 0)))(pe_params, X_pe, Y_pe)
raised = False
except TypeError as e:
raised = True
print("grad(vmap(loss)) -> TypeError:", str(e)[:90])
assert raised
# 3) Application A: DP-SGD-style clipping: clip each example's gradient to global L2 norm C, then average and add Gaussian noise
def per_example_norms(g):
# Flatten each leaf to (B, -1), sum the squares, then add up across leaves
return jnp.sqrt(sum(jnp.sum(l.reshape(l.shape[0], -1) ** 2, axis=1) for l in jax.tree.leaves(g)))
def clip_per_example(g, C):
scale = jnp.minimum(1.0, C / (per_example_norms(g) + 1e-12))
return jax.tree.map(lambda l: l * scale.reshape((-1,) + (1,) * (l.ndim - 1)), g)
C, sigma = 1.0, 0.5
g_clipped = clip_per_example(g_vmap, C)
norms_before, norms_after = per_example_norms(g_vmap), per_example_norms(g_clipped)
assert jnp.all(norms_after <= C + 1e-5) # After clipping, no example's norm exceeds C
assert jnp.allclose(norms_after, jnp.minimum(norms_before, C), atol=1e-4)
leaves, treedef = jax.tree.flatten(g_clipped)
noise_keys = jax.random.split(jax.random.key(7), len(leaves)) # An independent key for each leaf
dp_grad = jax.tree.unflatten(treedef, [l.mean(0) + sigma * C / B * jax.random.normal(k, l.shape[1:])
for l, k in zip(leaves, noise_keys)])
print(f"Max norm before clipping {norms_before.max():.2f} -> after clipping {norms_after.max():.2f}")
print("The DP gradient has the same shapes as the parameters:", jax.tree.map(jnp.shape, dp_grad) == jax.tree.map(jnp.shape, pe_params))
# 4) Application B: find outliers by per-example gradient norm (here, the 5 "mislabeled" examples)
norms = per_example_norms(g_vmap)
top5 = jnp.argsort(norms)[-5:]
print("5 examples with the largest gradient norms:", jnp.sort(top5), "| actually mislabeled:", bad_idx)
print("Median norm of outliers %.2f vs median norm of normal examples %.2f" % (
jnp.median(norms[bad_idx]), jnp.median(jnp.delete(norms, bad_idx))))
assert set(np.asarray(top5).tolist()) == set(np.asarray(bad_idx).tolist())
# 5) Application C: per-example loss weighting: gradient of the weighted loss == weighted sum of per-example gradients
weights = jnp.ones(B).at[bad_idx].set(0.0) # Example: give the outliers we found weight 0
weights = weights / weights.sum()
def weighted_loss(p, X, Y, w):
return jnp.sum(w * jax.vmap(mlp_loss, in_axes=(None, 0, 0))(p, X, Y))
g_weighted = jax.jit(jax.grad(weighted_loss))(pe_params, X_pe, Y_pe, weights)
g_weighted_from_per_ex = jax.tree.map(lambda l: jnp.tensordot(weights, l, axes=1), g_vmap)
assert tree_allclose(g_weighted, g_weighted_from_per_ex)
print("✔ grad(sum_i w_i * loss_i) == sum_i w_i * g_i")
print("Norm of the mean gradient before vs. after dropping outliers: %.3f -> %.3f" % (
per_example_norms(jax.tree.map(lambda l: l[None], g_mean_loss))[0],
per_example_norms(jax.tree.map(lambda l: l[None], g_weighted))[0]))
Memory considerations. The total size of the per-example gradients = batch size $B$ × number of parameters $P$ (4 bytes per element in float32), while an ordinary mean gradient is only $P$. For example, a model with 100 million parameters and batch size 32 needs $32 \times 10^8 \times 4\ \text{B} \approx 12.8\ \text{GB}$ for per-example gradients. Common ways to deal with this:
- When you only need per-example norms (e.g. DP-SGD clipping, outlier detection), compute in chunks with
lax.map(..., batch_size=k)and keep only the norms for each chunk, not the full gradients; - Compute the sum of clipped gradients chunk by chunk and accumulate;
- For large models, look into techniques like "ghost clipping" that never materialize per-example gradients.
# 6) Memory: per-example gradient size = B × P; compare compiled output size and temporary memory
n_elems = sum(l.size for l in jax.tree.leaves(g_vmap))
assert n_elems == B * P
print(f"Per-example gradients: {n_elems} elements = B({B}) × P({P}) ≈ {n_elems * 4 / 2**20:.2f} MiB; the mean gradient is only {P * 4 / 2**10:.2f} KiB")
for name, fn in [("jit(vmap(grad))", per_ex_grad), ("jit(grad(mean))", jax.jit(jax.grad(mean_mlp_loss)))]:
ma = fn.lower(pe_params, X_pe, Y_pe).compile().memory_analysis()
print(f"{name:16s} output {ma.output_size_in_bytes / 2**10:8.1f} KiB | temp {ma.temp_size_in_bytes / 2**10:8.1f} KiB")
# When you only need norms: compute in chunks with lax.map, return only norms per chunk; peak memory scales with chunk size
@jax.jit
def norms_chunked(p, X, Y):
one = lambda xy: per_example_norms(jax.tree.map(lambda l: l[None], jax.grad(mlp_loss)(p, xy[0], xy[1])))[0]
return lax.map(one, (X, Y), batch_size=32)
nc = norms_chunked(pe_params, X_pe, Y_pe)
assert jnp.allclose(nc, norms, rtol=1e-4, atol=1e-5)
ma = norms_chunked.lower(pe_params, X_pe, Y_pe).compile().memory_analysis()
print(f"lax.map in chunks, norms only: output {ma.output_size_in_bytes / 2**10:.1f} KiB | temp {ma.temp_size_in_bytes / 2**10:.1f} KiB, results match ✔")
🔍 Extra: vmap and random numbers, mapping over keys¶
🟢 Must-read
If you want each example to use different random numbers, first get n keys with jax.random.split(key, n), then pass the keys to vmap as a mapped argument. Never let all examples share the same key (in_axes=None), or every example gets exactly the same "random" numbers. For more key-management patterns (fold_in, scan carry, etc.), see Advanced 10 §5.
def noisy_sample(key, mu):
return mu + jax.random.normal(key, (2,))
keys = jax.random.split(jax.random.key(7), 4)
good = jax.vmap(noisy_sample, in_axes=(0, None))(keys, 0.0)
bad = jax.vmap(noisy_sample, in_axes=(None, 0))(jax.random.key(7), jnp.zeros(4))
print("A different key for each example:\n", good)
print("Wrong: all share the same key:\n", bad)
assert len(np.unique(np.asarray(good[:, 0]))) == 4 # all 4 examples differ
assert jnp.all(bad == bad[0]) # shared key: all 4 examples are identical
The top-level alias jax.tree_map, common in older tutorials, has been removed in JAX 0.12 (jax.tree_util.tree_map still works). This tutorial uses jax.tree.map throughout. For a complete old-to-new mapping, see Chapter 6, extra in 6.4 "PyTree API migration table".
# ❌ Old style (outdated): jax.tree_map was removed in JAX 0.12
try:
jax.tree_map(jnp.shape, g_vmap)
removed = False
except AttributeError as e:
removed = True
print("Measured AttributeError:", str(e)[:120])
assert removed
# ✅ Recommended: jax.tree.map
print(jax.tree.map(jnp.shape, g_vmap))
4.7 Performance considerations for vmap¶
vmap is a convenient tool for vectorizing functions, but it isn't magic, and it isn't necessarily faster than every alternative. It performs very well in many common cases, but you should keep its costs in both performance and memory in mind.
How vmap works (briefly). A Python for loop executes Python bytecode on every iteration. vmap instead transforms the function before execution (trace once, then apply batching rules), pushing the "loop" down into JAX so XLA can generate highly optimized parallel code for the whole batched operation (especially noticeable on GPUs/TPUs). That eliminates repeated Python calls and interpreter overhead.
vmap vs. manual vectorization. You have two options:
- Manual vectorization: rewrite the function with
jax.numpyso it handles batched arrays directly; vmap: write the function for a single example and letvmaphandle the batch dimension.
With jit, the two run at essentially the same speed, because vmap's batching rules produce ordinary batched operations. For the elementwise addition below, the two jaxprs are even identical character for character. For general functions, the generated program may differ from the hand-written one in details (for example, whether a transpose comes before or after a multiplication; an assert below shows one case), but the amount of computation is comparable. Only without jit does vmap add a little fixed overhead, because it redoes the transformation in Python on every call. The real value of vmap is for functions that are hard to vectorize by hand: complicated dimensions, deep nesting, or control flow written with jnp.where / lax.cond / lax.fori_loop and the like.
⚠️ Note: you often hear that "manual vectorization is usually as fast as vmap, or even slightly faster, because vmap has a small transformation overhead". That's only true without jit. In this example the two jaxprs are identical (checked with
assertbelow), and after jit they run at the same speed.
vmap and control flow. vmap can handle two kinds of control flow: static Python control flow (like a for with a fixed iteration count, which gets unrolled at trace time), and JAX control flow such as jnp.where, lax.cond and lax.while_loop. A Python if/else that depends on the values of batched data can't be vmapped and raises TracerBoolConversionError.
⚠️ Note: it's easy to believe that vmap can directly vectorize "functions with Python control flow (
if/else, elementwiseforloops)", generating batched code for the whole function, control flow included, from one trace. In fact a data-dependent Pythonifraises an error right away under vmap and must be rewritten withjnp.where/lax.cond(see the ❌/✅ comparison below).
vmap vs. Python loops. Calling a JAX function repeatedly from a Python loop has a lot of interpreter overhead, and JAX can't optimize across iterations. vmap expresses the computation as larger fused operations and is almost always much faster (two to four orders of magnitude faster on CPU in the measurements below).
Works best with jit. jax.jit(jax.vmap(f)) first applies the vectorizing transformation, then compiles the resulting batched function. XLA can fuse operations, optimize memory access patterns for the hardware, and parallelize efficiently along the batch dimension. Compilation also removes the overhead of launching many small computations from Python. The course includes a log-scale bar chart comparing the relative cost of "adding two vectors" written different ways; below we redraw it with numbers measured on this machine. Actual results depend on the complexity of the operation, the batch size and the hardware.
Memory use. vmap can speed things up a lot, but compared with a sequential loop its peak memory can be higher: JAX usually has to create arrays that hold the intermediate results for the whole batch at once. If f produces a large intermediate matrix, vmap(f) needs roughly batch_size × memory per intermediate matrix, while sequential execution needs only one at a time. Ways to cope:
- Reduce the batch size and process the data in chunks;
- Optimize the function itself to need less intermediate memory;
- Use
jax.lax.map(f, xs, batch_size=k): vectorized within a chunk, sequential across chunks, a compromise between memory and speed (measured at the end of this section).
⚠️ Note: when memory is tight, you don't need to fall back to a Python loop (with
jax.jit(f)inside).lax.map(..., batch_size=k)also keeps peak memory under control, runs at the same order of magnitude as vmap, and is far faster than a Python loop, so prefer it.
When not to use vmap:
- Simple operations that are already vectorized: writing them by hand may be simpler, with the same performance;
- Very tight memory: when the intermediate results for a whole batch don't fit, switch to chunking or
lax.map(batch_size=...); - Sequential dependencies between examples: if item $i$ depends on the result for item $i-1$,
vmapdoesn't apply; usejax.lax.scan.
To sum up: vmap is especially powerful combined with jit and is usually much faster than a Python loop, but keep an eye on memory and use profiling to compare different approaches.
Course code¶
Let's actually run it with the bench defined earlier (warm up first, then time with block_until_ready()).
batch_size = 1000
a = jnp.ones(batch_size)
b = jnp.arange(batch_size, dtype=jnp.float32)
# Manual vectorization
def manual_vectorized_add(x_batch, y_batch):
return jnp.add(x_batch, y_batch)
# Scalar function + vmap
def scalar_add(x, y):
return x + y
vmapped_add = jax.vmap(scalar_add)
# Python loop
def loop_add(x_batch, y_batch):
results = []
for i in range(len(x_batch)):
results.append(scalar_add(x_batch[i], y_batch[i]))
return jnp.stack(results)
# vmap + jit
jit_vmapped_add = jax.jit(vmapped_add)
assert jnp.allclose(manual_vectorized_add(a, b), vmapped_add(a, b))
assert jnp.allclose(loop_add(a, b), jit_vmapped_add(a, b))
timings = {
"Python loop": bench(loop_add, a, b, n=2),
"Manual (no jit)": bench(manual_vectorized_add, a, b),
"vmap (no jit)": bench(vmapped_add, a, b),
"jit(vmap)": bench(jit_vmapped_add, a, b),
}
for k_, v_ in timings.items():
print(f"{k_:18s} {v_:10.4f} ms")
# Corresponds to the bar chart in the course: relative time on a log scale (measured data)
import matplotlib.pyplot as plt
plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False
fig, ax = plt.subplots(figsize=(6, 3))
names, vals = list(timings.keys()), list(timings.values())
ax.barh(names, vals, color="#4C72B0")
ax.set_xscale("log")
ax.set_xlabel("Time per call (ms, log scale)")
ax.set_title("Adding two vectors of length 1000 (CPU)")
ax.invert_yaxis()
for i, v_ in enumerate(vals):
ax.text(v_, i, f" {v_:.3g}", va="center")
plt.tight_layout()
plt.show()
In this example, manual vectorization and vmap are the same program (in general they're equivalent programs):
jp_manual = jaxpr_str(jax.make_jaxpr(manual_vectorized_add)(a, b))
jp_vmap = jaxpr_str(jax.make_jaxpr(vmapped_add)(a, b))
print(jp_manual)
print(jp_vmap)
assert jp_manual == jp_vmap # Check: the two jaxprs are identical
# In general they aren't necessarily identical: for a linear layer with shared weights, hand-written X @ W.T and the vmap version put the transpose in different places, but the results agree
Wl, Xl = jnp.ones((3, 4)), jnp.arange(20.0).reshape(5, 4)
lin_vmap = jax.vmap(lambda W, x: W @ x, in_axes=(None, 0))
assert jaxpr_str(jax.make_jaxpr(lambda W, X: X @ W.T)(Wl, Xl)) != jaxpr_str(jax.make_jaxpr(lin_vmap)(Wl, Xl))
assert jnp.allclose(Xl @ Wl.T, lin_vmap(Wl, Xl))
jit_manual = jax.jit(manual_vectorized_add)
print(f"jit(manual): {bench(jit_manual, a, b, n=2000):.4f} ms | jit(vmap): {bench(jit_vmapped_add, a, b, n=2000):.4f} ms (microsecond-level differences are timing noise)")
vmap and data-dependent Python if:
# ❌ Wrong: expecting vmap to handle a data-dependent Python if/else directly
def piecewise_py(x):
if x > 0:
return x ** 2
return -x
try:
jax.vmap(piecewise_py)(jnp.array([-2.0, 3.0]))
raised = False
except jax.errors.TracerBoolConversionError as e:
raised = True
print("Measured", type(e).__name__, ":", str(e).splitlines()[0])
assert raised, "A Python if that depends on batched data should raise under vmap"
# ✅ Recommended: express data-dependent branches with jnp.where (or lax.cond); static for loops can be used as is
def piecewise_ok(x):
return jnp.where(x > 0, x ** 2, -x)
out = jax.vmap(piecewise_ok)(jnp.array([-2.0, 3.0]))
assert jnp.array_equal(out, jnp.array([2.0, 9.0]))
print("jnp.where version:", out)
def static_loop(x): # The loop count is static: fine, it's unrolled at trace time
for _ in range(3):
x = jnp.sin(x)
return x
n_sin = jaxpr_str(jax.make_jaxpr(jax.vmap(static_loop))(jnp.ones(2))).count("sin")
assert n_sin == 3
print("Static for loop:", jax.vmap(static_loop)(jnp.array([0.5, 1.0])), "| unrolled into", n_sin, "sin ops in the jaxpr")
🔍 Extra: measuring again with a heavier function¶
⚪ Optional: read this if you want to see a timing comparison with more computation.
The addition above is so simple that almost all the time is call overhead. Below we switch to the forward pass of a small two-layer network and compare a Python loop, vmap and jit(vmap).
k1, k2 = jax.random.split(jax.random.key(8))
W1 = jax.random.normal(k1, (128, 128)) / jnp.sqrt(128.0)
Xp = jax.random.normal(k2, (2000, 128))
def mlp_single(W, x): # Forward pass of a small network for one example; the weights are passed in as an argument (see the extra below for why)
h = jnp.tanh(W @ x)
return jnp.sum(jnp.tanh(W @ h) ** 2)
loop_fn = lambda W, X: jnp.stack([mlp_single(W, x) for x in X])
vmap_fn = jax.vmap(mlp_single, in_axes=(None, 0))
jit_vmap_fn = jax.jit(jax.vmap(mlp_single, in_axes=(None, 0)))
# Cross-check against a NumPy float64 reference (and confirm we don't hit the YNN bug described below)
W1n, Xpn = np.asarray(W1, np.float64), np.asarray(Xp[:50], np.float64)
ref_np = np.sum(np.tanh(np.tanh(Xpn @ W1n.T) @ W1n.T) ** 2, axis=1)
assert np.allclose(np.asarray(loop_fn(W1, Xp[:50])), ref_np, rtol=1e-4, atol=1e-4)
assert np.allclose(np.asarray(jit_vmap_fn(W1, Xp[:50])), ref_np, rtol=1e-4, atol=1e-4)
t_loop = bench(loop_fn, W1, Xp, n=1)
t_vmap = bench(vmap_fn, W1, Xp)
t_jv = bench(jit_vmap_fn, W1, Xp)
print(f"Python loop: {t_loop:9.2f} ms")
print(f"vmap: {t_vmap:9.3f} ms")
print(f"jit(vmap): {t_jv:9.3f} ms (about {t_loop / t_jv:.0f}x faster than the loop)")
🔍 Extra (measured in this environment): pass large arrays as arguments; don't let a jit function capture them in a closure¶
🟢 Must-read
The timing cell above deliberately passes the weights W as an argument, for two reasons:
- General reason: an array captured by a
jitfunction's closure is embedded as a constant in the compiled program, making the program bigger and compilation slower. And if the closure variable is reassigned later, the cached version keeps using the old value. Model parameters should be passed in as function arguments. - A known bug in this environment: the CPU backend of jaxlib 0.12.0.1 compiles "closure-constant matrix
W @ xfollowed by a sum" into code that silently returns wrong results. A common pattern likejax.jit(jax.vmap(lambda x: jnp.sum(jnp.tanh(W @ x))))(withWas a closure variable) hits it exactly. For the symptoms, a subprocess reproduction and workarounds, see "Known issues in this environment" in 00 Overview; for the root-cause analysis, see Advanced Chapter 2 §2.3.
All other examples in this chapter have been checked: whenever a large array is used inside
jit, it is passed in as a function argument (likeW,MandWmin this section, and the MLP parameters in the 4.6 special topic). Key results are additionally cross-checked withassertagainst a NumPy float64 reference or an analytic solution.
🔍 Extra: under batching, lax.cond becomes select, and both branches run¶
🟢 Must-read
If the predicate of a lax.cond depends on mapped data, different examples in the batch may take different branches, and a vectorized program can't send "some examples down A and some down B". So vmap rewrites the cond as: compute both branches on the whole batch, then pick results elementwise with select_n.
Consequences:
- The cost is the sum of both branches; in this situation
condsaves no computation; - If the unselected branch produces
nan/inffor some inputs (say,logof a negative number), the forward result is unaffected. The backward pass is unaffected too, and this is where it differs fromjnp.where: a batchedcondfirst wraps the inputs of unselected examples instop_gradientbefore handing them to the branch (seestop_gradient+select_nin the jaxpr below), so infinite derivatives there can't flow back to the input. A hand-writtenjnp.wherehas no such protection and gives NaN gradients (see Chapter 3, 3.8); - More subtly, every operand used by the branches is broadcast to the batch dimension (verified below).
Only when the predicate doesn't depend on mapped data (say, it comes from an argument with in_axes=None) does cond stay a real conditional branch.
def g_cond(x):
return lax.cond(x > 0, jnp.sin, jnp.cos, x)
print("===== Single example: a real cond =====")
print(jax.make_jaxpr(g_cond)(1.0))
print("===== After vmap (batched predicate): both branches computed, then select_n picks =====")
jp_c = jaxpr_str(jax.make_jaxpr(jax.vmap(g_cond))(jnp.arange(3.0)))
print(jp_c)
assert "cond[" not in jp_c and "select_n" in jp_c and "sin" in jp_c and "cos" in jp_c # both branches are there
# When the predicate comes from a shared argument (in_axes=None), cond stays a cond
g_flag = lambda x, flag: lax.cond(flag, jnp.sin, jnp.cos, x)
jaxpr_txt = jaxpr_str(jax.make_jaxpr(jax.vmap(g_flag, in_axes=(0, None)))(jnp.arange(3.0), True))
assert "cond[" in jaxpr_txt
print("When the predicate doesn't depend on batched data, the jaxpr still contains cond ✔")
# Unlike jnp.where: a batched cond adds stop_gradient to the inputs of the unselected branch, so the unselected log branch doesn't make the gradient nan
xs_g = jnp.array([-1.0, 0.0, 2.0])
g_c = jax.vmap(jax.grad(lambda x: lax.cond(x > 0, jnp.log, lambda x: 0.0 * x, x)))(xs_g)
g_w = jax.vmap(jax.grad(lambda x: jnp.where(x > 0, jnp.log(x), 0.0)))(xs_g)
print("Gradient of batched cond:", g_c, "| gradient of jnp.where:", g_w)
assert not jnp.any(jnp.isnan(g_c)) and bool(jnp.isnan(g_w[1]))
# Check with timing and memory: the expensive branch is never selected, yet you still pay for it
M = jax.random.normal(jax.random.key(9), (256, 256)) / 16.0
def expensive(M, x):
return lax.fori_loop(0, 30, lambda i, v: jnp.tanh(M @ v), x)
def cheap(M, x):
return x * 2.0
def branchy(M, x):
return lax.cond(jnp.sum(x) > 1e9, expensive, cheap, M, x) # the predicate is false for every example
def branchy_where(M, x): # explicitly "compute both, then pick"
return jnp.where(jnp.sum(x) > 1e9, expensive(M, x), cheap(M, x))
xs_c = jax.random.normal(jax.random.key(10), (512, 256))
fs = {name: jax.jit(jax.vmap(fn, in_axes=(None, 0)))
for name, fn in [("cheap", cheap), ("expensive", expensive), ("cond", branchy), ("where", branchy_where)]}
assert jnp.allclose(fs["cond"](M, xs_c), fs["cheap"](M, xs_c))
mems = {}
for name, fn in fs.items():
mem = mems[name] = fn.lower(M, xs_c).compile().memory_analysis().temp_size_in_bytes / 2**20
print(f"vmap({name:9s}) time {bench(fn, M, xs_c, n=5):8.3f} ms | temp memory {mem:8.2f} MiB")
# Why: a batched cond broadcasts every operand (including the shared M) to the batch dimension
jp = jaxpr_str(jax.make_jaxpr(jax.vmap(branchy, in_axes=(None, 0)))(M, xs_c))
assert "f32[512,256,256]" in jp and mems["cond"] > 50 * mems["expensive"]
print("✔ The jaxpr contains a broadcast M (f32[512,256,256]); cond's temp memory is %.0fx that of expensive" % (mems["cond"] / mems["expensive"]))
Conclusion (results on CPU): a batched
conddoesn't just "run both branches". It also broadcasts every operand (including the shared weightsMwithin_axes=None) into arrays with a batch dimension, so the shared matrix-vector product turns into 512 independent matrix products. Temporary memory grows from about 1 MiB to about 128 MiB, and the run time is an order of magnitude higher than just running the expensive branch. The equivalentjnp.whereversion also computes both branches, but the shared weights aren't broadcast, so the cost is roughly "expensive branch + cheap branch". Rule of thumb: inside vmap, if the predicate depends on batched data and the branches use large shared parameters, preferjnp.where(or move the shared computation outside thecond). Don't assumecondwill save computation.
🔍 Extra: a batched while_loop keeps running until the "slowest example" finishes¶
⚪ Optional: read this if examples in a batch need very different numbers of iterations.
After vmap, the loop condition of a lax.while_loop becomes "any example in the batch is still running" (an OR reduction over the per-example conditions), and the loop body uses select to keep the state of finished examples unchanged. So every example's result is correct, but the number of iterations is set by the slowest example.
Below we demonstrate this with Collatz sequence step counts, and find reduce ... or (the condition reduction) and select (keeping finished examples' state) in the compiled StableHLO.
def collatz_steps(n):
def cond_fn(s):
return s[0] != 1
def body_fn(s):
v, k = s
return jnp.where(v % 2 == 0, v // 2, 3 * v + 1), k + 1
return lax.while_loop(cond_fn, body_fn, (n, 0))[1]
ns = jnp.array([1, 2, 3, 6, 7, 27])
steps = jax.vmap(collatz_steps)(ns)
print("Steps for each number:", steps) # 27 needs 111 steps, so the whole batch loops 111 times
assert steps.tolist() == [0, 1, 7, 8, 16, 111]
hlo = jax.jit(jax.vmap(collatz_steps)).lower(ns).as_text()
assert "stablehlo.or" in hlo and "stablehlo.select" in hlo
print("✔ The condition has an OR reduction across the batch dimension; the body uses select to hold finished examples")
🔍 Extra: memory and speed trade-offs between jax.lax.map and vmap¶
⚪ Optional: read this when vmap runs out of memory (OOM).
lax.map(f, xs) has the same semantics as vmap(f)(xs), but underneath it uses lax.scan and runs the examples sequentially, one at a time (a loop compiled into the program, not a Python loop).
vmap: the most parallelism and the fastest, but peak memory grows linearly with the batch size;lax.map: peak memory is only as much as for one example, but it's fully serial and slow;lax.map(f, xs, batch_size=k)(supported in JAX 0.12): a compromise. It processes k examples at a time with vmap and runs the chunks sequentially with scan. If the number of examples isn't divisible by k, the leftover examples are handled separately (visible in the jaxpr below).
Below we look at temporary memory with the compiled memory_analysis() and compare speed.
k1, k2 = jax.random.split(jax.random.key(11))
Wm = jax.random.normal(k1, (512, 512)) / jnp.sqrt(512.0)
Xm = jax.random.normal(k2, (4096, 512))
def heavy(W, x): # One example produces a (512,) intermediate, (4096, 512) when batched; W is passed in as an argument
return jnp.tanh(W @ x) @ W
variants = {
"vmap": jax.jit(jax.vmap(heavy, in_axes=(None, 0))),
"lax.map": jax.jit(lambda W, xs: lax.map(lambda x: heavy(W, x), xs)),
"lax.map(batch_size=256)": jax.jit(lambda W, xs: lax.map(lambda x: heavy(W, x), xs, batch_size=256)),
}
Wmn = np.asarray(Wm, np.float64)
ref = np.tanh(np.asarray(Xm, np.float64) @ Wmn.T) @ Wmn # NumPy float64 reference: heavy(x) = tanh(W x) @ W
temps = {}
for name, fn in variants.items():
mem = fn.lower(Wm, Xm).compile().memory_analysis()
temp = temps[name] = mem.temp_size_in_bytes / 1024
assert np.allclose(np.asarray(fn(Wm, Xm)), ref, atol=1e-3)
print(f"{name:26s} temp memory {temp:9.1f} KiB | time {bench(fn, Wm, Xm, n=5):8.2f} ms | matches NumPy reference ✔")
assert temps["lax.map"] < temps["lax.map(batch_size=256)"] < temps["vmap"]
# How it handles a number of examples not divisible by batch_size
print(jax.make_jaxpr(lambda xs: lax.map(lambda x: x ** 2, xs, batch_size=4))(jnp.arange(10.0)))
How to read this:
vmapuses the most temporary memory (intermediates for the whole batch);lax.mapthe least, but it's much slower;batch_sizesits in between, running at the same order of magnitude asvmapand much faster than per-examplelax.map. The exact numbers depend on the machine; these are CPU results. On GPU/TPU, the advantage fromvmap's parallelism is usually even more pronounced. Rule of thumb: start withvmap; when you run out of memory (OOM), switch tolax.map(..., batch_size=k)and tune k to balance memory and speed.
4.8 Hands-on: vectorizing functions¶
This section practices jax.vmap with a few examples: from vectorizing the simplest function, to controlling the mapping with in_axes, to combining it with jit and grad. The core idea: write clear single-example logic first, then apply it to batches with vmap, differentiate with grad, and compile with jit.
# Environment check
import jax
import jax.numpy as jnp
import timeit
print("Available devices:", jax.devices())
Course example 1: basic vectorization¶
First write a function on scalars and process a batch of inputs with a manual loop, then switch to vmap. A Python loop works, but it's usually slow, especially with large batches, complex functions, or on accelerators. By default vmap maps over axis 0 of the input, and axis 0 of the output corresponds to the mapped dimension.
def process_scalar(x):
return x * 2.0 + 1.0
scalar_input = 5.0
scalar_output = process_scalar(scalar_input)
print(f"Input: {scalar_input}, output: {scalar_output}")
batch_inputs = jnp.arange(1.0, 6.0) # Array([1., 2., 3., 4., 5.])
print("Batch input:", batch_inputs)
# Manual loop
manual_outputs = []
for x in batch_inputs:
manual_outputs.append(process_scalar(x))
manual_outputs = jnp.array(manual_outputs)
print("Manual Loop Outputs:", manual_outputs)
# vmap
vectorized_process = jax.vmap(process_scalar)
vmap_outputs = vectorized_process(batch_inputs)
print("vmap Outputs:", vmap_outputs)
print("Input shape:", batch_inputs.shape, "| output shape:", vmap_outputs.shape)
assert jnp.array_equal(manual_outputs, vmap_outputs)
Course example 2: controlling the mapping with in_axes¶
in_axes is a tuple that says which axis to map over for each argument; None means the argument is broadcast, not mapped. Below, each vector gets its own scale factor, while the shift is shared by all examples. In the output shape (3, 3), the first 3 is the batch dimension added by vmap and the second 3 comes from the shape of a single vector. If the batch dimension is the last axis, you can write in_axes=(..., -1, ...).
def scale_and_shift(vector, scale, shift):
return vector * scale + shift
single_vector = jnp.array([1.0, 2.0, 3.0])
single_scale = 2.0
single_shift = 10.0
print("Single input:", scale_and_shift(single_vector, single_scale, single_shift))
batch_vectors = jnp.array([[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0],
[7.0, 8.0, 9.0]]) # shape (3, 3)
batch_scales = jnp.array([0.5, 1.0, 1.5]) # shape (3,)
fixed_shift = 100.0 # all examples use the same shift
vectorized_scale_shift = jax.vmap(scale_and_shift, in_axes=(0, 0, None))
batch_outputs = vectorized_scale_shift(batch_vectors, batch_scales, fixed_shift)
print("batch_vectors:", batch_vectors.shape, "| batch_scales:", batch_scales.shape)
print(batch_outputs)
print("Output shape:", batch_outputs.shape)
assert jnp.allclose(batch_outputs, batch_vectors * batch_scales[:, None] + fixed_shift)
Course example 3: combining vmap, jit and grad¶
The course uses simple_activation(params, x) = tanh(params * x) to demonstrate combining the three transformations: first jit on its own, then the gradient with respect to params, then processing a batch of inputs in two orders (vmap then jit, jit then vmap), and finally vectorizing the gradient function and timing it.
Note that jax.grad can only differentiate functions with scalar output, while simple_activation(1.5, single_x) returns a vector of shape (5,). So there are two meaningful ways to "take the gradient with respect to params":
- Sum the output:
grad(lambda p, x: jnp.sum(simple_activation(p, x)))gives the total gradient of all elements with respect toparams; - Elementwise gradients: the single-example function handles only a scalar
x, and vmap maps it over each element (two layers of vmap for batched data). This matches the meaning of an "elementwise activation function", and it's exactly where vmap shines. The two satisfy "total gradient = sum of elementwise gradients".
As for the two orders of composition, they give the same result; jit(vmap(f)) is recommended, for the reasons in 4.6. When timing, warm up first, or you'll mostly be measuring compilation.
⚠️ Note: you can't call
jax.grad(simple_activation)(single_param, single_x)directly on a vector output, andvmap(grad_activation, in_axes=(None, 0))on batched data doesn't work either. Both raiseTypeError: Gradient only defined for scalar-output functions. The ❌/✅ comparison below gives two fixes and checks them withassertagainst the analytic result $\partial_p \tanh(p x) = x\,(1-\tanh^2(p x))$.
⚠️ Note: it is not the case that "jit first, then vmap is better". The two orders give the same result and essentially the same compiled kernels, and putting
jiton the outside has the least Python overhead (consistent with the conclusion in 4.6). Also, timing only the first call ofjit_vectorized_gradmostly measures compilation, so the timing below warms up first.
def simple_activation(params, x):
return jnp.tanh(params * x)
single_x = jnp.linspace(-2.0, 2.0, 5)
single_param = 1.5
single_output = simple_activation(single_param, single_x)
print("Single output:", single_output)
# JIT
jit_activation = jax.jit(simple_activation)
jit_output = jit_activation(single_param, single_x)
assert jnp.allclose(single_output, jit_output)
print("JIT output matches the original output.")
# ❌ Wrong: calling grad directly on a vector output
grad_activation = jax.grad(simple_activation, argnums=0)
try:
single_grad = grad_activation(single_param, single_x)
raised = False
except TypeError as e:
raised = True
print("Measured TypeError:", str(e)[:110])
assert raised
# ✅ Recommended
# Fix 1: sum the output to get a scalar -> total gradient
grad_activation_sum = jax.grad(lambda p, x: jnp.sum(simple_activation(p, x)), argnums=0)
g_sum = grad_activation_sum(single_param, single_x)
# Fix 2: elementwise gradients: the single-example function handles only a scalar x (so the output is a scalar and grad is valid), then vmap maps it over each element
grad_activation_elem = jax.grad(simple_activation, argnums=0)
per_elem_grad = jax.vmap(grad_activation_elem, in_axes=(None, 0))(single_param, single_x)
analytic = single_x * (1 - jnp.tanh(single_param * single_x) ** 2) # d/dp tanh(p x) = x (1 - tanh(p x)^2)
print("Fix 1, total gradient:", g_sum, "| Fix 2, elementwise gradients:", per_elem_grad)
assert jnp.allclose(per_elem_grad, analytic, atol=1e-6) # Fix 2 matches the analytic result
assert jnp.allclose(g_sum, analytic.sum(), atol=1e-6) # Fix 1 matches the sum of the analytic result
assert jnp.allclose(per_elem_grad.sum(), g_sum, atol=1e-6) # total gradient = sum of elementwise gradients
print("✔ Both fixes match the analytic result")
# Batched input
batch_x = jnp.stack([single_x, single_x * 0.5, single_x * -1.0]) # shape (3, 5)
print("Batched input shape:", batch_x.shape)
# Option 1: vmap first, then jit (recommended)
vectorized_activation = jax.vmap(simple_activation, in_axes=(None, 0))
jit_vectorized_activation = jax.jit(vectorized_activation)
batch_output_1 = jit_vectorized_activation(single_param, batch_x)
# Option 2: jit first, then vmap
vectorized_jit_activation = jax.vmap(jit_activation, in_axes=(None, 0))
batch_output_2 = vectorized_jit_activation(single_param, batch_x)
assert batch_output_1.shape == batch_output_2.shape == (3, 5)
assert jnp.allclose(batch_output_1, batch_output_2)
print("The outputs of both composition orders match.")
# ❌ Wrong: vmap(grad_activation) over batched data; each row is still a vector output
try:
jax.vmap(grad_activation, in_axes=(None, 0))(single_param, batch_x)
raised = False
except TypeError as e:
raised = True
print("Measured TypeError:", str(e)[:110])
assert raised
# ✅ Recommended: two layers of vmap; the outer one loops over rows, the inner one over each scalar element in a row
vectorized_grad = jax.vmap(jax.vmap(grad_activation_elem, in_axes=(None, 0)), in_axes=(None, 0))
batch_grads = vectorized_grad(single_param, batch_x)
print("Batched gradient shape:", batch_grads.shape)
print(batch_grads)
assert jnp.allclose(batch_grads, batch_x * (1 - jnp.tanh(single_param * batch_x) ** 2), atol=1e-6)
jit_vectorized_grad = jax.jit(vectorized_grad)
start = timeit.default_timer()
batch_grads_jit = jit_vectorized_grad(single_param, batch_x).block_until_ready()
print(f"First call (includes compilation): {timeit.default_timer() - start:.6f}s")
print(f"Average after warm-up: {bench(jit_vectorized_grad, single_param, batch_x) / 1e3:.6f}s")
assert jnp.allclose(batch_grads, batch_grads_jit)
Course exercise: pairwise Euclidean distances¶
Given two sets of 2-D points:
points_a = jnp.array([[0., 0.], [1., 1.]]) # shape (2, 2)
points_b = jnp.array([[2., 2.], [3., 3.], [4., 4.]]) # shape (3, 2)
First write a function that computes the Euclidean distance between two points, then use vmap (nested if you like) to build a distance matrix whose (i, j) entry is the distance between points_a[i] and points_b[j]. The expected shape is (2, 3). Hint: think carefully about the in_axes of each vmap layer.
Self-check interface: euclidean_distance(p1, p2) returns the distance between two points (a scalar); compute_pairwise_distances(arr_a, arr_b) returns an (n, m) distance matrix for two point sets of shapes (n, d) and (m, d), and should work for any n, m, d.
# ✍️ Your answer: Course exercise: pairwise Euclidean distances
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Course exercise: pairwise Euclidean distances
try:
euclidean_distance, compute_pairwise_distances
except NameError as e:
print(f"⏳ Course exercise: not answered yet: {e}")
else:
assert abs(float(euclidean_distance(jnp.array([0., 0.]), jnp.array([3., 4.]))) - 5.0) < 1e-5, \
"euclidean_distance((0, 0), (3, 4)) should be 5"
_D = np.asarray(compute_pairwise_distances(jnp.array([[0., 0.], [1., 1.]]),
jnp.array([[2., 2.], [3., 3.], [4., 4.]])))
assert _D.shape == (2, 3), f"For the data in the problem the distance matrix should have shape (2, 3); got {_D.shape}. Check the in_axes of each vmap layer"
# New data: different numbers of points, dimension 3; reference computed with NumPy broadcasting
_ka, _kb = jax.random.split(jax.random.key(2024))
_A = np.asarray(jax.random.normal(_ka, (5, 3)))
_B = np.asarray(jax.random.normal(_kb, (7, 3)))
_D2 = np.asarray(compute_pairwise_distances(jnp.asarray(_A), jnp.asarray(_B)))
assert _D2.shape == (5, 7), f"5 points against 7 points should give shape (5, 7); got {_D2.shape} (rows correspond to arr_a, columns to arr_b)"
_ref = np.sqrt(((_A[:, None, :] - _B[None, :, :]) ** 2).sum(-1))
assert np.allclose(_D2, _ref, atol=1e-5), "Entry (i, j) should be the Euclidean distance between arr_a[i] and arr_b[j]"
print("✅ Course exercise: self-check passed")
Note: if you later need gradients of this distance matrix and the two point sets share a point, the gradient of
sqrtat 0 will benan; use the "doublewhere" trick (see Chapter 3, 3.8).
Below are extra exercises, a bit harder.
Extra exercise 1(a): batched MLP prediction (shared dict parameters)¶
Given a single-example function predict(params, x) (a two-layer network, x of shape (d,)), use vmap to write the batched version batched_predict(params, X), where X has shape (B, d) and the parameters are shared. Compare it with hand-written batched matrix operations to confirm the results agree, then compile it with jit and time it.
Self-check interface: predict(params, x) is the two-layer ReLU network W2 @ relu(W1 @ x + b1) + b2, where params is a dict with W1 (h, d), b1 (h,), W2 (o, h), b2 (o,); batched_predict(params, X) returns (B, o) for X of shape (B, d).
# ✍️ Your answer: Extra exercise 1(a)
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise 1(a)
try:
batched_predict, predict
except NameError as e:
print(f"⏳ Extra exercise 1(a): not answered yet: {e}")
else:
_k = jax.random.split(jax.random.key(2025), 5)
_p = {"W1": jax.random.normal(_k[0], (5, 7)), "b1": jax.random.normal(_k[1], (5,)),
"W2": jax.random.normal(_k[2], (3, 5)), "b2": jax.random.normal(_k[3], (3,))}
_X = jax.random.normal(_k[4], (11, 7))
_out = np.asarray(batched_predict(_p, _X))
assert _out.shape == (11, 3), f"The output shape should be (B, o) = (11, 3); got {_out.shape}"
_pn = {k: np.asarray(v, np.float64) for k, v in _p.items()}
_H = np.maximum(np.asarray(_X, np.float64) @ _pn["W1"].T + _pn["b1"], 0)
assert np.allclose(_out, _H @ _pn["W2"].T + _pn["b2"], atol=1e-4), \
"The batched result doesn't match a hand-written NumPy two-layer network (parameters should be shared across all examples)"
_loop = np.stack([np.asarray(predict(_p, x)) for x in _X])
assert np.allclose(_out, _loop, atol=1e-4), "The batched result doesn't match calling predict on each example"
assert np.asarray(batched_predict(_p, _X[:1])).shape == (1, 3), "With B = 1 the output shape should be (1, 3)"
print("✅ Extra exercise 1(a): self-check passed")
Extra exercise 1(b): in_axes / out_axes when examples are stored as columns¶
The data matrix Xc has shape (d, B) (each column is an example). Without transposing the data, use in_axes to have predict process each column directly, and use out_axes to make the output shape (d_out, B).
Self-check interface: reuse predict and params from 1(a); Xc is data of shape (d, B) (one example per column), and out_cols is the batched prediction on it, with shape (d_out, B).
# ✍️ Your answer: Extra exercise 1(b)
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise 1(b)
try:
out_cols, Xc, predict, params
except NameError as e:
print(f"⏳ Extra exercise 1(b): not answered yet: {e}")
else:
_Xc, _o = np.asarray(Xc), np.asarray(out_cols)
assert _Xc.shape[0] == np.asarray(params["W1"]).shape[1], \
f"Xc should be (d, B), with axis 0 the feature dimension d = {np.asarray(params['W1']).shape[1]}; got {_Xc.shape}"
_ref = np.stack([np.asarray(predict(params, jnp.asarray(_Xc[:, j]))) for j in range(_Xc.shape[1])], axis=1)
assert _o.shape == _ref.shape, f"The output shape should be (d_out, B) = {_ref.shape}; got {_o.shape}. Check out_axes"
assert np.allclose(_o, _ref, atol=1e-4), "Column j of out_cols should equal predict(params, Xc[:, j])"
print("✅ Extra exercise 1(b): self-check passed")
Extra exercise 2: DP-SGD-style per-example gradient clipping¶
Use a linear regression model ($y = w^\top x + b$, with true $w = (1, -2, 0.5)$, $b = 0.5$):
- Compute per-example gradients with
vmap(grad); - Compute the global L2 norm of each example's gradient (combining the gradients of
wandb); - Scale down gradients whose norm exceeds
Cso their norm is exactlyC, then average; - Run 200 steps of gradient descent with it and check that
wconverges close totrue_w.
(Real DP-SGD also adds Gaussian noise; here we only practice the vectorization part.)
Self-check interface: loss_single(theta, x, y) returns the single-example squared error $(w^\top x+b-y)^2$, with theta = {"w": (d,), "b": scalar}; clipped_mean_grad(theta, X, Y, C) returns (clipped mean gradient, per-example gradient norms), where the first is a dict with the same structure as theta and the second has shape (N,).
# ✍️ Your answer: Extra exercise 2
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise 2
try:
clipped_mean_grad, loss_single
except NameError as e:
print(f"⏳ Extra exercise 2: not answered yet: {e}")
else:
_k1, _k2 = jax.random.split(jax.random.key(2026))
_X = jax.random.normal(_k1, (40, 4)) * 3
_Y = jax.random.normal(_k2, (40,)) * 3
_th = {"w": jnp.array([0.3, -0.2, 0.1, 0.5]), "b": jnp.array(-0.4)}
# Reference: closed-form per-example gradients of linear regression, d/dw = 2r·x, d/db = 2r (float64)
_Xn, _Yn = np.asarray(_X, np.float64), np.asarray(_Y, np.float64)
_r = _Xn @ np.asarray(_th["w"], np.float64) + float(_th["b"]) - _Yn
assert np.isclose(float(loss_single(_th, _X[0], _Y[0])), _r[0] ** 2, rtol=1e-4), "loss_single should be the single-example squared error"
_gw, _gb = 2 * _r[:, None] * _Xn, 2 * _r
_norms = np.sqrt((_gw ** 2).sum(1) + _gb ** 2)
for _C in (0.5, 5.0, 1e6): # 1e6: no clipping happens, so this should reduce to the ordinary mean gradient
_s = np.minimum(1.0, _C / _norms)
_g, _nm = clipped_mean_grad(_th, _X, _Y, _C)
assert np.asarray(_nm).shape == (40,), f"Per-example norms should have shape (N,) = (40,); got {np.asarray(_nm).shape}"
assert np.allclose(_nm, _norms, rtol=1e-3), "Per-example norms are wrong: compute the L2 norm over the gradients of w and b together"
assert np.allclose(_g["w"], (_s[:, None] * _gw).mean(0), rtol=1e-3, atol=1e-4), f"With C = {_C}, the clipped mean gradient of w is wrong"
assert np.isclose(float(_g["b"]), (_s * _gb).mean(), rtol=1e-3, atol=1e-4), f"With C = {_C}, the clipped mean gradient of b is wrong"
print("✅ Extra exercise 2: self-check passed")
Extra exercise 3: simulating many random walks in parallel with vmap (mapping over keys)¶
Write a single-path function walk(key, n_steps): each step goes +1 or -1 with equal probability, and it returns the final position. Use vmap to simulate 10000 paths at once, and check that the mean final position is about 0 and the variance is about n_steps.
Self-check interface: walk(key, n_steps) returns the final position (a scalar) of an n_steps-step ±1 random walk; the self-check calls it with fresh keys and n_steps = 50.
# ✍️ Your answer: Extra exercise 3
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise 3
try:
walk
except NameError as e:
print(f"⏳ Extra exercise 3: not answered yet: {e}")
else:
_keys = jax.random.split(jax.random.key(2027), 20000)
_f = np.asarray(jax.jit(jax.vmap(lambda k: walk(k, 50)))(_keys)).astype(np.float64)
assert _f.shape == (20000,), f"walk should return a scalar (the final position), so the batched shape should be (20000,); got {_f.shape}"
assert np.all(np.abs(_f) <= 50) and np.all((_f + 50) % 2 == 0), \
"The end point of a 50-step ±1 walk should be in [-50, 50] and even; check that each step is only +1 or -1"
assert abs(_f.mean()) < 0.3, f"The mean end point should be about 0; got {_f.mean():.3f}"
assert abs(_f.var() - 50) < 3, f"The variance of the end point should be about n_steps = 50; got {_f.var():.2f}"
assert float(walk(_keys[0], 50)) == float(walk(_keys[0], 50)), "The same key should give the same result"
assert len(np.unique(_f[:200])) > 5, "Different keys should give different paths; check that you actually use the key passed in"
print("✅ Extra exercise 3: self-check passed")
Extra exercise 4: batched Jacobians and jnp.vectorize¶
For $f(x) = \big(\sin x_0 \cdot x_1,\ x_0^2 + x_1\big)$:
- Use
vmap(jax.jacfwd(f))to compute the2×2Jacobian at each point of a(B, 2)batch, and check the first point by hand; - Do the same with
jnp.vectorize(..., signature='(n)->(m,n)').jnp.vectorizeis built on vmap and uses a NumPy broadcasting (gufunc) signature to handle any number of leading batch dimensions automatically.
Self-check interface: f_vec(x) returns the $f(x)$ above (shape (2,)) for x of shape (2,); jac_vectorized(X) returns (..., 2, 2) Jacobians for input of shape (..., 2), with any number of leading batch dimensions.
# ✍️ Your answer: Extra exercise 4
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Extra exercise 4
try:
jac_vectorized, f_vec
except NameError as e:
print(f"⏳ Extra exercise 4: not answered yet: {e}")
else:
assert np.allclose(f_vec(jnp.array([0.5, 2.0])), [np.sin(0.5) * 2.0, 0.25 + 2.0], atol=1e-6), \
"f_vec doesn't match the formula in the problem, (sin x0 · x1, x0² + x1)"
_x = np.asarray(jax.random.normal(jax.random.key(2028), (3, 4, 2)))
_J = np.asarray(jac_vectorized(jnp.asarray(_x)))
assert _J.shape == (3, 4, 2, 2), f"For input (3, 4, 2) the output should be (3, 4, 2, 2); got {_J.shape}. Check the signature"
_ref = np.empty((3, 4, 2, 2))
_ref[..., 0, 0] = np.cos(_x[..., 0]) * _x[..., 1]
_ref[..., 0, 1] = np.sin(_x[..., 0])
_ref[..., 1, 0] = 2 * _x[..., 0]
_ref[..., 1, 1] = 1.0
assert np.allclose(_J, _ref, atol=1e-5), "The Jacobian doesn't match the closed form (row i holds the partial derivatives of output i with respect to each input)"
print("✅ Extra exercise 4: self-check passed")
4.9 Chapter 4 Quiz¶
These 8 questions cover this chapter and help you check how well you've got it. All answers have been verified in code.
1. Which statement about jax.vmap(f)(xs) is correct?
- A. It calls
fone item at a time in a Pythonforloop at run time - B. Semantically it's equivalent to
jnp.stack([f(x) for x in xs]); it's implemented by applying batching rules to the traced program - C. It automatically runs
fin parallel on multiple threads - D. It only works for elementwise functions
2. In f(W, x), W is a shared weight and x is a batch of examples (batch dimension on axis 0). You should use:
- A.
in_axes=(0, 0)· B.in_axes=(None, 0)· C.in_axes=(0, None)· D.out_axes=None
3. In f(a, b), a.shape == (5, 3) and b.shape == (3, 5). You want to pair each row of a with each column of b. You should use:
- A.
in_axes=(0, 0)· B.in_axes=(1, 0)· C.in_axes=(0, 1)· D.in_axes=(None, 1)
4. With X.shape == (M, d) and Y.shape == (N, d), what is the output shape of vmap(vmap(f, in_axes=(None, 0)), in_axes=(0, None))(X, Y) (where f returns a scalar)?
- A.
(M,)· B.(N, M)· C.(M, N)· D.(M, N, d)
5. A single-example loss loss(params, x, y) returns a scalar. Which expression gives per-example gradients?
- A.
jax.grad(jax.vmap(loss, in_axes=(None, 0, 0))) - B.
jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)) - C.
jax.vmap(jax.grad(loss), in_axes=(0, 0, 0)) - D.
jax.jit(jax.grad(loss))
6. Under vmap, what happens if the predicate of a lax.cond depends on mapped data?
- A. It raises; vmap doesn't support
lax.cond - B. Each example runs only its own branch
- C. Both branches run on the whole batch, and select picks the results elementwise
- D. Only the branch for the first example runs
7. What does jax.vmap(lambda a, b: a + b)(jnp.ones(3), jnp.ones(4)) give?
- A. A matrix of shape
(3, 4)· B. Shape(4,); the shorter array is padded automatically - C. Raises
ValueError: mapped axis sizes don't match · D. Shape(3,); the extra elements are truncated
8. When the batch is huge and vmap runs out of memory, the best alternative is:
- A. Go back to a Python for loop
- B. Process in chunks with
jax.lax.map(f, xs, batch_size=k) - C. Remove
jit - D. Set every
in_axestoNone
# Verify some of the answers in code
# Question 3
a3, b3 = jnp.ones((5, 3)), jnp.ones((3, 5))
q3 = jax.vmap(jnp.dot, in_axes=(0, 1))(a3, b3).shape; assert q3 == (5,); print("Q3:", q3)
# Question 4
fq = lambda x, y: jnp.dot(x, y)
q4 = jax.vmap(jax.vmap(fq, in_axes=(None, 0)), in_axes=(0, None))(jnp.ones((2, 3)), jnp.ones((7, 3))).shape
assert q4 == (2, 7); print("Q4:", q4)
# Question 5: A raises, B gives per-example gradients
try:
jax.grad(jax.vmap(mlp_loss, in_axes=(None, 0, 0)))(pe_params, X_pe, Y_pe)
except TypeError:
print("Q5: A raises TypeError (the output isn't a scalar)")
print("Q5: shapes for B", jax.tree.map(jnp.shape, jax.vmap(jax.grad(mlp_loss), in_axes=(None, 0, 0))(pe_params, X_pe, Y_pe)))
# Question 7
try:
jax.vmap(lambda a, b: a + b)(jnp.ones(3), jnp.ones(4))
except ValueError as e:
print("Q7: ValueError ->", str(e).splitlines()[0])
Answers and explanations
- B. vmap is a function transformation: it traces
f, applies a batching rule to each primitive and produces a batched program, with no run-time loop (you can see this in the jaxpr comparison in 4.2). It isn't limited to elementwise functions either. - B.
Nonemeans not mapped (shared), and0means mapped over axis 0. - C.
ais mapped over axis 0 (rows) andbover axis 1 (columns); both have length 5, so they can be paired. - C. The outer batch dimension (over X, size M) comes first and the inner one (over Y, size N) second.
- B. Shared parameters use
None, and the data is mapped over axis 0. In A, the output ofvmap(loss)is a vector, sogradraises; C would require params to have a batch dimension too; D is just a single-example gradient. - C. A batched
condis rewritten asselect_nand both branches are computed (4.7 has the jaxpr and timing checks). It only stays a realcondwhen the predicate doesn't depend on batched data. - C. Mapped axes must have the same size; vmap doesn't do NumPy-style broadcasting.
- B.
lax.mapruns sequentially with scan, andbatch_sizekeeps each chunk vectorized, trading off memory and speed. A Python loop is the slowest, and removingjitdoesn't reduce the memory used by intermediates.
Chapter summary¶
- Vectorization: replace per-example loops with whole-array operations to avoid Python overhead and use hardware parallelism; manual vectorization makes it easy to get the batch dimension wrong.
jax.vmap(f): write only the single-example function and get the batched version automatically. Semantically it equals "call one by one, then stack"; it's implemented as a program transformation (batching rules), with no run-time loop.in_axes: an integer is the axis to map over (may be negative),Nonemeans shared; it can be a tuple or a pytree matching the argument structure (such as a dict). Keyword arguments are always mapped over axis 0.out_axes: controls which axis the output's batch dimension goes on;Noneonly works for outputs that don't vary across examples (or withsum_match=True, new in 0.12).axis_size: gives the batch size when no input is mapped.- Several batched arguments: arguments are paired one to one; mapped axis sizes must match and there's no automatic broadcasting.
- Nested vmap: each layer handles one batch dimension; use it for pairwise operations (distance matrices, kernel matrices). The outer batch dimension comes first in the output.
- Composing transformations:
jit(vmap(f))is the recommended order (vmap(jit(f))gives the same kernels but more Python overhead);gradneeds a scalar output. - Per-example gradients:
jit(vmap(grad(loss), in_axes=(None, 0, 0))), equivalent to the Jacobian of the per-example losses with respect to the parameters; their mean equals the gradient of the mean loss; useful for DP-SGD clipping, outlier detection and example weighting; memory is batch size × number of parameters, and when you only need norms, compute in chunks withlax.map(batch_size=...). - Random numbers: use
jax.random.key, and use each key only once;splitinto n keys first, then map over the keys. - Environment note: the CPU backend of jaxlib 0.12.0.1 has a YNN fusion bug, so pass large arrays to jit functions as arguments instead of capturing them in a closure (see 00 Overview).
- Performance: vmap is about as fast as hand-written batched code, and better with jit; the cost is peak memory that grows with batch size. When memory runs short, use
lax.map(..., batch_size=k). - Control flow: under batching,
condbecomesselect(both branches run), andwhile_loopruns until the slowest example finishes; a Pythonifcan't depend on batched data.
Cheat sheet¶
| API / pattern | What it does |
|---|---|
jax.vmap(f) | Map f over axis 0 of every positional argument |
jax.vmap(f, in_axes=(None, 0)) | 1st argument shared, 2nd mapped over axis 0 |
in_axes=1 / in_axes=-1 | Map over axis 1 / the last axis |
in_axes=({"w": None, "b": 0}, 0) | Use a pytree (dict) to set the mapping for different fields inside an argument |
jax.vmap(f, out_axes=1) | Put the output's batch dimension on axis 1 |
out_axes=None | Output without a batch dimension (must not vary across examples) |
jax.vmap(f, axis_size=n) | Give the batch size explicitly (required when all inputs are None) |
jax.vmap(f, out_axes=None, sum_match=True) | Sum an out_axes=None output along the batch dimension (JAX 0.12) |
jax.vmap(f, axis_name="i") + lax.psum/pmean(x, "i") | Collective operations over the batch dimension (e.g. normalizing within a batch) |
vmap(vmap(f, (None, 0)), (0, None)) | Pairwise operation, output (M, N) |
jax.jit(jax.vmap(f)) | Compiled batched function (recommended) |
jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)) | Per-example gradients |
jax.vmap(sample)(jax.random.split(key, n)) | Independent random numbers for each example |
jax.lax.map(f, xs, batch_size=k) | Run in sequential chunks to save memory |
jnp.vectorize(f, signature="(n)->(m,n)") | Built on vmap; supports any number of leading batch dimensions via a gufunc signature |
jnp.einsum("bnk,bkm->bnm", A, B) / A @ B | Batched matrix multiplication (equivalent to vmap(jnp.matmul)) |
jax.make_jaxpr(jax.vmap(f))(xs) | Show the program after the vmap transformation |
# Extra demo for the cheat sheet: axis_name + collective operations (centering within a batch)
center = jax.vmap(lambda x: x - lax.pmean(x, axis_name="batch"), axis_name="batch")
print(center(jnp.arange(4.0))) # [-1.5 -0.5 0.5 1.5]
Version notes¶
Below is a summary of the usage points and conclusions in this chapter that need attention under JAX 0.12. Each one is checked with assert at the corresponding place in the text.
| No. | Section | Key point | Where |
|---|---|---|---|
| 1 | 4.2 | "print runs only once" under vmap is caused by tracing, not by vmap optimizing execution: Python print runs once at trace time, while jax.debug.print runs once per example | 4.2 Introducing jax.vmap |
| 2 | 4.3 | The length of an in_axes tuple must match the number of positional arguments at call time; keyword arguments are always mapped over axis 0 and the tuple doesn't cover them | 4.3 Mapping over specific arguments (in_axes, out_axes) |
| 3 | 4.4 | The default in_axes=0 maps all positional arguments, not just the first; a shared scalar needs None, or you get ValueError | 4.4 Handling multiple batched arguments |
| 4 | 4.4 | Column scaling with vmap(scale_column, in_axes=(1, 0)) outputs (4, 3) (the batch dimension defaults to output axis 0); to get the 3×4 layout you must add out_axes=1 | 4.4 Handling multiple batched arguments |
| 5 | 4.5 | In the nested vmap for pairwise matvec, the inner (None, 0) loops over vectors and the outer (0, None) over matrices, giving (3, 4, 2) | 4.5 Nested vmap |
| 6 | 4.6 | Create keys with jax.random.key(seed) (typed keys); jax.random.PRNGKey(seed) is the old style | 4.6 Combining vmap with jit and grad |
| 7 | 4.6 | Using the same key for the uniform x and then the normal noise makes them highly correlated (correlation about 0.98, rank correlation 1); after jax.random.split the correlation is about 0 | 4.6 Combining vmap with jit and grad |
| 8 | 4.6 | vmap(jit(f)) pushes the batch dimension inside the jit, so its kernels are essentially the same as jit(vmap(f)); the only difference is Python-side overhead; jit(vmap(f)) is recommended | 4.6 Combining vmap with jit and grad |
| 9 | 4.6 | Per-example gradients special topic: theory, cross-check against a loop/jacrev, timing, the mean property, DP-SGD clipping, outlier detection, example weighting, memory of batch size × number of parameters | 4.6 Combining vmap with jit and grad |
| 10 | 4.6 | jax.tree_map was removed in 0.12 (raises AttributeError); use jax.tree.map | 4.6 Combining vmap with jit and grad |
| 11 | 4.7 | Manual vectorization is slightly faster than vmap only without jit; in this example the two jaxprs are identical (in general they're equivalent programs), and after jit they run at the same speed | 4.7 Performance considerations for vmap |
| 12 | 4.7 | A data-dependent Python if raises TracerBoolConversionError under vmap; use jnp.where / lax.cond; a static for gets unrolled and can be used as is | 4.7 Performance considerations for vmap |
| 13 | 4.7 | When memory is tight, prefer lax.map(f, xs, batch_size=k) over falling back to a Python loop: memory stays under control and speed is the same order of magnitude as vmap | 4.7 Performance considerations for vmap |
| 14 | 4.7 | Costs of control flow under batching: lax.cond becomes select, both branches run, and shared operands are broadcast to the batch dimension (about 128x more memory); while_loop runs until the slowest example finishes | 4.7 Performance considerations for vmap |
| 15 | 4.8 | In example 3, calling jax.grad directly on a vector output and vmap(grad) directly on batched data both raise TypeError; fix by "summing the output" or "two layers of elementwise vmap", cross-checked against the analytic result | 4.8 Hands-on: vectorizing functions |
| 16 | 4.8 | The two composition orders of jit and vmap give the same result; jit(vmap(f)) is recommended, consistent with 4.6 | 4.8 Hands-on: vectorizing functions |
| 17 | 4.8 | When timing jit_vectorized_grad, the first call is mostly compilation; warm up first, then average over several runs | 4.8 Hands-on: vectorizing functions |
| 18 | 4.3 | Usage and demos of pytree in_axes, axis_size, out_axes=None and sum_match (new in 0.12) | 4.3 Mapping over specific arguments (in_axes, out_axes) |
| 19 | 4.7 | jaxlib 0.12.0.1 CPU: a constant matrix captured in a jit function's closure, multiplied with data and then summed, is computed wrongly by YNN fusion; pass the parameters as function arguments, or set XLA_FLAGS=--xla_cpu_experimental_ynn_fusion_type=-reduce (subprocess reproduction and asserts in "Known issues in this environment" in 00 Overview) | 4.7 Performance considerations for vmap |
Appendix: reference solutions¶
First write your own answer in the "✍️ Your answer" cell in the main text and check it with the "🧪 Self-check". Come here to compare once the self-check passes, or if you're really stuck. Running the solutions below in order will overwrite any functions with the same names you defined in your answer cells.
Course exercise: pairwise Euclidean distances¶
The problem is in the main text under "Course exercise: pairwise Euclidean distances".
# Reference solution (the course's solution)
def euclidean_distance(p1, p2):
return jnp.sqrt(jnp.sum((p1 - p2)**2))
points_a = jnp.array([[0., 0.], [1., 1.]])
points_b = jnp.array([[2., 2.], [3., 3.], [4., 4.]])
# Hint: first map over b only -> one a against all b, shape (3,)
vmap_over_b = jax.vmap(euclidean_distance, in_axes=(None, 0))
print("points_a[0] to all of b:", vmap_over_b(points_a[0], points_b))
@jax.jit
def compute_pairwise_distances(arr_a, arr_b):
vmap_dist_b = jax.vmap(euclidean_distance, in_axes=(None, 0)) # inner: loop over b
vmap_dist_a_b = jax.vmap(vmap_dist_b, in_axes=(0, None)) # outer: loop over a
return vmap_dist_a_b(arr_a, arr_b)
distance_matrix = compute_pairwise_distances(points_a, points_b)
print("points_a:\n", points_a)
print("points_b:\n", points_b)
print("Distance matrix:\n", distance_matrix)
print("Shape:", distance_matrix.shape)
# Expected output given by the course: [[2.828427, 4.2426405, 5.656854], [1.4142135, 2.828427, 4.2426405]]
assert distance_matrix.shape == (2, 3)
assert jnp.allclose(distance_matrix, jnp.array([[2.828427, 4.2426405, 5.656854], [1.4142135, 2.828427, 4.2426405]]))
print("✔ Matches the expected output")
Extra exercise 1(a)¶
The problem is in the main text under "Extra exercise 1(a): batched MLP prediction (shared dict parameters)".
# Reference solution
def init_mlp(key, d_in, d_hidden, d_out):
k1, k2 = jax.random.split(key)
return {
"W1": jax.random.normal(k1, (d_hidden, d_in)) / jnp.sqrt(d_in),
"b1": jnp.zeros(d_hidden),
"W2": jax.random.normal(k2, (d_out, d_hidden)) / jnp.sqrt(d_hidden),
"b2": jnp.zeros(d_out),
}
def predict(params, x): # Single example: x has shape (d,)
h = jax.nn.relu(params["W1"] @ x + params["b1"])
return params["W2"] @ h + params["b2"]
params = init_mlp(jax.random.key(12), 16, 32, 4)
Xe = jax.random.normal(jax.random.key(13), (256, 16))
batched_predict = jax.jit(jax.vmap(predict, in_axes=(None, 0)))
def manual_batched(params, X): # Hand-written batched version, for comparison
H = jax.nn.relu(X @ params["W1"].T + params["b1"])
return H @ params["W2"].T + params["b2"]
print(batched_predict(params, Xe).shape)
assert jnp.allclose(batched_predict(params, Xe), manual_batched(params, Xe), atol=1e-5)
print("✔ Matches the hand-written version")
print("jit(vmap) time: %.3f ms" % bench(batched_predict, params, Xe))
Extra exercise 1(b)¶
The problem is in the main text under "Extra exercise 1(b): in_axes / out_axes when examples are stored as columns".
# Reference solution
Xc = Xe.T # (16, 256): examples stored as columns
out_cols = jax.vmap(predict, in_axes=(None, 1), out_axes=1)(params, Xc)
print("Output shape:", out_cols.shape) # (4, 256)
assert out_cols.shape == (4, 256) and jnp.allclose(out_cols, batched_predict(params, Xe).T, atol=1e-5)
print("✔ Matches the row-wise version")
Extra exercise 2¶
The problem is in the main text under "Extra exercise 2: DP-SGD-style per-example gradient clipping".
# Reference solution
k1, k2 = jax.random.split(jax.random.key(15))
Xd = jax.random.normal(k1, (256, 3))
true_w = jnp.array([1.0, -2.0, 0.5])
Yd = Xd @ true_w + 0.5 + 0.1 * jax.random.normal(k2, (256,))
def loss_single(theta, x, y): # Single-example loss of linear regression (scalar)
pred = jnp.dot(theta["w"], x) + theta["b"]
return (pred - y) ** 2
def clipped_mean_grad(theta, X, Y, C=1.0):
g = jax.vmap(jax.grad(loss_single), in_axes=(None, 0, 0))(theta, X, Y)
# Global norm per example: sum of squares over the non-batch dimensions of each leaf, then add up across leaves
sq = sum(jnp.sum(leaf.reshape(leaf.shape[0], -1) ** 2, axis=1) for leaf in jax.tree.leaves(g))
norms = jnp.sqrt(sq) # shape (N,)
scale = jnp.minimum(1.0, C / (norms + 1e-12)) # scale factor for each example
g = jax.tree.map(lambda t: t * scale.reshape((-1,) + (1,) * (t.ndim - 1)), g)
return jax.tree.map(lambda t: jnp.mean(t, axis=0), g), norms
@jax.jit
def step(theta, X, Y, lr=0.1):
g, norms = clipped_mean_grad(theta, X, Y)
return jax.tree.map(lambda p, gi: p - lr * gi, theta, g), jnp.max(norms)
th = {"w": jnp.zeros(3), "b": jnp.array(0.0)}
for i in range(200):
th, max_norm = step(th, Xd, Yd)
if i in (0, 199):
print(f"Step {i}: max per-example gradient norm = {max_norm:.3f}")
print("Learned w:", th["w"], "| true w:", true_w)
assert jnp.allclose(th["w"], true_w, atol=0.05) and abs(float(th["b"]) - 0.5) < 0.05
print("Learned b:", th["b"], "| true b: 0.5")
Extra exercise 3¶
The problem is in the main text under "Extra exercise 3: simulating many random walks in parallel with vmap (mapping over keys)".
# Reference solution
def walk(key, n_steps=100):
steps = jax.random.choice(key, jnp.array([-1, 1]), shape=(n_steps,))
return jnp.sum(steps)
keys = jax.random.split(jax.random.key(16), 10000)
finals = jax.jit(jax.vmap(walk))(keys)
print("Mean:", jnp.mean(finals.astype(jnp.float32)), "(theory: 0)")
print("Variance:", jnp.var(finals.astype(jnp.float32)), "(theory: 100)")
assert abs(float(jnp.mean(finals))) < 0.5 and abs(float(jnp.var(finals.astype(jnp.float32))) - 100) < 5
Extra exercise 4¶
The problem is in the main text under "Extra exercise 4: batched Jacobians and jnp.vectorize".
# Reference solution
def f_vec(x):
return jnp.array([jnp.sin(x[0]) * x[1], x[0] ** 2 + x[1]])
pts = jax.random.normal(jax.random.key(17), (5, 2))
J = jax.vmap(jax.jacfwd(f_vec))(pts)
print("Batched Jacobian shape:", J.shape) # (5, 2, 2)
x0 = pts[0]
J0_manual = jnp.array([[jnp.cos(x0[0]) * x0[1], jnp.sin(x0[0])],
[2 * x0[0], 1.0]])
assert jnp.allclose(J[0], J0_manual, atol=1e-6)
print("✔ The first point matches the hand derivation")
jac_vectorized = jnp.vectorize(jax.jacfwd(f_vec), signature="(n)->(m,n)")
assert jnp.allclose(jac_vectorized(pts), J)
print("✔ jnp.vectorize gives the same result")
print("Any number of leading batch dimensions:", jac_vectorized(pts.reshape(5, 1, 2)).shape)