Neural Networks & LLM Components · Easy

PyTorch Pairwise Distances

Use singleton dimensions and broadcasting to compute every row-pair distance.

Task

Implement pairwise_distances(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor. Compute the Euclidean distance between every row of x and every row of y using PyTorch broadcasting.

Requirements

x has shape (N, F) and y has shape (M, F).
Return a tensor with shape (N, M).
Use broadcasting with singleton dimensions instead of explicit Python loops.
Reduce over the feature dimension only.
Preserve the input dtype and device in the returned tensor.

Example

pairwise_distances(
    torch.tensor([[0.0, 0.0], [3.0, 4.0]]),
    torch.tensor([[0.0, 0.0], [6.0, 8.0]]),
)
# tensor([[ 0., 10.],
#         [ 5.,  5.]])