Skip to content

← Broadcasting step 2 of 6

Medium Primitives

Every pair at once

points is (n, d): n points in d dimensions. You want the (n, n) matrix of Euclidean distances between every pair.

The nested loop is n^2 iterations of Python. The broadcast is one expression, and the trick behind it is worth learning once because it generalises to every all-pairs computation you will ever write.

Make the axes disagree on purpose

points[:, None, :]      # (n, 1, d)
points[None, :, :]      # (1, n, d)

Subtract them. Broadcasting aligns from the right:

(n, 1, d)
(1, n, d)
---------
(n, n, d)

Entry [i, j, :] is points[i] - points[j], because axis 0 varies over the first copy and axis 1 over the second. Square, sum over the last axis, take the square root, and you have the distance matrix.

diff = points[:, None, :] - points[None, :, :]     # (n, n, d)
(diff ** 2).sum(-1).sqrt()                         # (n, n)

Read the memory before you ship it

That intermediate is n * n * d elements. For 10,000 points in 128 dimensions it is 1.28e10 floats, which is 51 GB, and the expression that produced it looks like one line of arithmetic.

This is the characteristic failure of broadcasting: it is not slow, it is enormous, and the size never appears in the source. torch.cdist exists precisely to compute this without materialising the difference tensor, and for large inputs it is the right call. Write the broadcast when you need to understand it or when n is small; reach for cdist when it is not.

Your task

def pairwise(points: torch.Tensor) -> torch.Tensor

Return the (n, n) matrix of Euclidean distances, using at most eight dispatched operations, and without torch.cdist.

The starter is the double loop.