Skip to content

← JAX Vectorization & Control Flow step 4 of 25

Hard Primitives

Nested vmap (vmap-of-vmap)

Jump to the task

Why this matters

A pairwise distance table compares every point in one set with every point in another. If X contains N points and Y contains M points, the table has N * M distances.

Start with a function for one pair. One vmap compares a fixed point against all of Y. A second vmap repeats that operation across X, producing one row per point in X.

Worked mini-example

import jax
import jax.numpy as jnp

X = jnp.array([[0.0, 0.0], [1.0, 1.0]])
Y = jnp.array([[0.0, 0.0], [3.0, 4.0], [1.0, 0.0]])

def distance(x, y):
    return jnp.linalg.norm(x - y)

one_row = jax.vmap(distance, in_axes=(None, 0))
all_rows = jax.vmap(one_row, in_axes=(0, None))
result = all_rows(X, Y)
# Approximately [[0, 5, 1], [1.414, 3.606, 1]], shape (2, 3)

The inner map shares one x and maps over the rows of Y. Its result has shape (M,). The outer map supplies each row of X while sharing the entire Y, producing shape (N, M).

Common pitfalls

  • One map with in_axes=(0, 0) pairs corresponding rows. It does not calculate all combinations, and it requires equal mapped lengths.
  • The outer map calls the already-mapped function. Read its arguments at that level: one point from X, plus all of Y.
  • Mapping X inside and Y outside with default output axes instead puts Y on the rows and X on the columns.
  • None means shared within that map, not static for JIT.
  • Vectorization does not remove the storage cost of the result. The full distance table still contains N * M entries.

Problem

Implement pairwise_distances(X, Y) using nested jax.vmap, without loops or jnp.einsum.

Inputs have shapes (N, d) and (M, d). Return shape (N, M), with out[i, j] = jnp.linalg.norm(X[i] - Y[j]).

Return Euclidean distances, including the square root, not squared distances. Identical points have distance zero. If either point set is empty, keep the corresponding output dimension zero.

Inner map: one row of distances

Share one point x of shape [2] and map over Y of shape [3, 2]. Each distance is a scalar ([]), so the three distances form one row [3].

Loading visualization…

Outer map: stack the rows

Map over X of shape [2, 2], sharing all of Y. The inner map returns a row [3] for each point in X; stacking the two rows gives the full [2, 3] result. These panels show the two stages separately; their controls are independent.

Loading visualization…