Skip to content

← Broadcasting step 2 of 6

Medium Primitives

Compute pairwise Euclidean distances

Jump to the task

Given n points in d dimensions, compute the Euclidean distance between every pair. The input has shape (n, d) and the result has shape (n, n):

distance[i, j] = sqrt(sum((points[i, k] - points[j, k]) ** 2 over k))

Give each point index its own axis

Insert singleton axes so one operand varies over the first point index and the other varies over the second:

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

At [i, j, :], the difference tensor contains points[i] - points[j]. Reduce its feature axis to turn each difference vector into one distance:

import torch

points = torch.tensor([[0., 0.], [3., 4.]])
diff = points[:, None, :] - points[None, :, :]
distances = diff.square().sum(dim=-1).sqrt()
print(list(diff.shape))  # [2, 2, 2]
print(distances.tolist())  # [[0.0, 5.0], [5.0, 0.0]]

The diagonal is zero because each point is compared with itself. Different rows can also have zero distance if they contain the same point.

Count the intermediate elements

The inserted-axis views share input storage, but the subtraction creates n * n * d elements in this eager expression. With 10,000 points and 128 features, that is 12.8 billion elements. At four bytes per float32 value, the difference tensor alone needs 51.2 GB in decimal units. Squaring it and other operations can add to peak memory use.

The final (n, n) matrix is also quadratic in the number of points. For 10,000 points, its float32 values need another 400 MB. Removing the large difference tensor does not remove the cost of storing every distance.

torch.cdist provides a dedicated pairwise-distance interface. Other approaches can process blocks or avoid storing the full matrix when only nearest neighbors are needed. The choice depends on the required output, memory budget and numerical accuracy; a short expression is not evidence that a calculation is cheap.

Numerical limits still apply

The square-and-sum formula can overflow for large magnitudes and lose precision for poorly represented inputs. The exercise uses small values to focus on axis alignment. A different implementation may use another calculation order and produce slightly different floating-point results.

Your task

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

Return the (n, n) Euclidean distance matrix for the nonempty 2-D floating-point input. Do not call torch.cdist. Use at most eight dispatched operations on the harness inputs.

The harness checks distances rounded to four decimal places and counts PyTorch dispatches. It does not measure GPU kernel count or runtime. Replace the starter's nested Python loop with tensor operations.