Batched Pairwise Euclidean Distance Matrices in PyTorch

Hello again. In the previous lesson, you distinguished points from the displacements between them and saw why translating an entire point cloud leaves those displacements unchanged. Today you will turn each displacement’s length into an entry of a pairwise distance matrix. By the end, you will be able to compute such matrices for batches of 3D point clouds in PyTorch, check their shapes and values, and choose between two useful implementations.

Plan on about 40 minutes, including the readings and code.

From points to a distance matrix

Suppose a cloud $X$ contains $N$ points and a cloud $Y$ contains $M$ points. Each row stores one point in 3D. Their Euclidean distance matrix $D$ has shape $N\times M$, with

A row tells you how far one point in $X$ is from every point in $Y$. If $X=Y$, the matrix is square and contains all within-cloud distances, including zeros on its diagonal.

For a batch, let x have shape (B, N, 3) and y have shape (B, M, 3). We want D to have shape (B, N, M): D[b, i, j] compares points x[b, i] and y[b, j]. It does not compare points from different batch elements. The same calculation also works for feature vectors of dimension other than three, provided both inputs have the same final dimension.

PyTorch provides this operation directly. Read its API description with particular attention to which dimensions represent batches, points, and coordinates.

Compute distances directly

Here is a small batch whose answers we can check by inspection. In the first pair of clouds, $(3,0,0)$ is five units from $(0,4,0)$. The second pair is the same geometry shifted to a different location.

import torch

dtype = torch.float64

x0 = torch.tensor(
    [[0., 0., 0.],
     [3., 0., 0.],
     [0., 4., 0.]],
    dtype=dtype,
)
y0 = torch.tensor(
    [[0., 0., 0.],
     [0., 4., 0.]],
    dtype=dtype,
)

shift = torch.tensor([10., -5., 2.], dtype=dtype)
x = torch.stack((x0, x0 + shift))  # (B=2, N=3, 3)
y = torch.stack((y0, y0 + shift))  # (B=2, M=2, 3)

D = torch.cdist(x, y, p=2)         # (2, 3, 2)
expected = torch.tensor(
    [[0., 4.],
     [3., 5.],
     [4., 0.]],
    dtype=dtype,
)

assert D.shape == (2, 3, 2)
assert torch.allclose(D[0], expected)
assert torch.allclose(D[1], expected)
print(D[0])

Both batches produce the same distances because corresponding points in x and y received the same translation. If only one cloud were shifted, its distances to the other cloud would generally change.

You can also compute the distances by explicitly forming every pairwise displacement. PyTorch’s None indexing inserts axes of length one:

differences = x[:, :, None, :] - y[:, None, :, :]
# x side: (B, N, 1, 3)
# y side: (B, 1, M, 3)
# result: (B, N, M, 3)

D_direct = torch.linalg.vector_norm(differences, dim=-1)
assert torch.allclose(D_direct, D)

Broadcasting replicates values conceptually across the point-pair axes; you do not write Python loops over points or batches. The final dim=-1 reduces the three coordinate differences to one length per pair. This direct form closely matches the definition, making it especially useful for checking shape logic.

Derive the matrix form

There is another way to calculate distances without first materializing a displacement vector for every pair. Expand the squared Euclidean distance using the dot product:

The first term depends only on $i$, the second only on $j$, and the dot products for all pairs can be obtained by matrix multiplication. The following reading develops precisely this broadcasting pattern. Its displayed code may render imperfectly, so concentrate on the explanation of the three terms; the complete batched implementation follows here.

For (B, N, 3) and (B, M, 3) inputs, that identity becomes:

x_norm2 = x.square().sum(dim=-1, keepdim=True)    # (B, N, 1)
y_norm2 = y.square().sum(dim=-1).unsqueeze(-2)    # (B, 1, M)
dot_products = x @ y.transpose(-1, -2)            # (B, N, M)

D2 = x_norm2 + y_norm2 - 2 * dot_products         # squared distances
D_from_matrix = D2.clamp_min(0).sqrt()            # Euclidean distances

assert D2.shape == (2, 3, 2)
assert torch.allclose(D2, D.square())
assert torch.allclose(D_from_matrix, D)

The @ performs one matrix multiplication per batch element. Its result has one entry for each point pair; the two arrays of squared norms broadcast across its columns and rows. Notice the distinction between D2 and D: a squared distance of $25$ means a Euclidean distance of $5$, not $25$.

Why clamp before the square root? The mathematical expression cannot be negative, but floating-point subtraction can produce a tiny negative result near zero. Clamping prevents an invalid square root. The matrix formula can also lose precision when coordinates are very large but their differences are small: it subtracts large, nearly equal numbers. For an initial implementation, torch.cdist or the direct difference formula is usually the clearer choice; the matrix form is valuable for understanding the computation and can avoid storing a (B, N, M, 3) displacement tensor. Neither form avoids the (B, N, M) output itself. That quadratic growth matters for large point clouds.

Check geometry, not just shapes

When comparing a cloud to itself, distances should be symmetric and diagonal entries should be zero, up to floating-point tolerance:

self_distances = torch.cdist(x, x, p=2)  # (B, N, N)

assert torch.allclose(
    self_distances,
    self_distances.transpose(-1, -2),
)
assert torch.allclose(
    self_distances.diagonal(dim1=-2, dim2=-1),
    torch.zeros((2, 3), dtype=dtype),
    atol=1e-12,
)

For two different clouds, the result need not be square or symmetric. A more general check follows from the previous lesson: if the same translation $\mathbf{t}$ is applied to both inputs, every displacement is unchanged because

A common rotation also preserves distances. For an orthogonal rotation matrix $Q$, the reason is $Q^{\top }Q=I$:

Scaling does not generally preserve distances. These facts give you stronger tests than checking a few hard-coded entries: they test whether an implementation respects the geometry it is meant to compute.

A batched pairwise Euclidean distance matrix has one entry per point pair within each batch element. torch.cdist(x, y, p=2) computes it directly; broadcasting pairwise differences reveals its definition; and the dot-product expansion computes squared distances through batched matrix multiplication. Keep track of shapes and of the difference between distances and squared distances, and use symmetry and shared-transform checks to catch mistakes.

Next, you will study eigendecomposition of symmetric matrices and why their eigenvectors can form an orthonormal basis—another linear-algebra tool that will later help describe geometric data.