1 / 1

Pairwise Distances with Nested vmap

JAX

Medium

Compute the matrix of squared Euclidean distances and the matrix of cosine similarities between every point of a and every point of b.

For point $a_i$ (row i of a) and point $b_j$ (row j of b),

S_{ij} = \sum_{d=1}^{D} \left(a_{id} - b_{jd}\right)^2 , \qquad C_{ij} = \frac{a_i \cdot b_j}{\max\left(\lVert a_i \rVert \, \lVert b_j \rVert,\ 10^{-8}\right)} ,

where $a_i \cdot b_j$ is the dot product and $\lVert a_i \rVert$ is the Euclidean norm, the square root of $a_i \cdot a_i$ . The floor of $10^{-8}$ makes $C_{ij}$ equal to 0 whenever either point is the zero vector. The input a has shape (N, D) and b has shape (M, D).

Return a dict with two float32 JAX arrays of shape (N, M), sq_dist holding $S$ and cosine holding $C$ , where row i and column j hold the values for $a_i$ and $b_j$ .

Examples

Input: a = [[1,0],[0,2]], b = [[1,0],[0,1],[-1,-1]]

Output: {"sq_dist":[[0.0,2.0,5.0],[5.0,1.0,10.0]],"cosine":[[1.0,0.0,-0.707107],[0.0,1.0,-0.707107]]}

Explanation: Point (0, 2) of a and point (-1, -1) of b differ by (1, 3), so their squared distance is 1 + 9 = 10. Their dot product is -2 and their norms are 2 and 1.414214, so their cosine is -2 / 2.828427 = -0.707107.

Input: a = [[0,0,0],[1,2,2]], b = [[2,4,4],[0,0,0]]

Output: {"sq_dist":[[36.0,0.0],[9.0,9.0]],"cosine":[[0.0,0.0],[1.0,0.0]]}

Input: a = [[3,4],[-3,-4],[0.5,0]], b = [[3,4]]

Output: {"sq_dist":[[0.0],[100.0],[22.25]],"cosine":[[1.0],[-1.0],[0.6]]}

Hint 1

jax.vmap(f, in_axes=(None, 0)) maps over rows of b with one row of a.

Hint 2

jax.vmap(jax.vmap(f, in_axes=(None, 0)), in_axes=(0, None)) gives an (N, M) result.

Hint 3

If f returns (d1, d2), vmap returns two batched arrays.

Requirements

Write a function for one pair of points and build the (N, M) grid with one jax.vmap nested inside another, one mapping over the rows of a and one over the rows of b, with no Python loops.
The function must also give the same result when it is wrapped in jax.jit.

Constraints

N and M are between 1 and 5, and D is between 1 and 6.
a and b share one dtype, float32 or int32, and inputs from the Run panel are float32.
All values are finite, points may be the zero vector, and a and b may contain identical points.

Try Similar Problems

Batched Attention with vmapMedium Model Ensemble with in_axesEasy Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Global Norm Gradient ClippingMedium AdamW from ScratchMedium Mixed Precision with Loss ScalingMedium