pyagc.utils.pairwise_squared_distance

pairwise_squared_distance(x: Tensor, y: Tensor) → Tensor

Compute pairwise squared Euclidean distances between two sets of vectors.

Efficiently computes the squared \(L_2\) distance between all pairs of vectors from two sets using the identity:

\[\| \mathbf{x}_i - \mathbf{y}_j \|_2^2 = \| \mathbf{x}_i \|_2^2 - 2 \mathbf{x}_i^\top \mathbf{y}_j + \| \mathbf{y}_j \|_2^2\]

where \(\mathbf{x}_i\) and \(\mathbf{y}_j\) are the \(i\)-th and \(j\)-th vectors in sets \(x\) and \(y\) respectively.

This vectorized implementation is more efficient than naive nested loops and is commonly used in clustering algorithms (e.g., K-Means) and nearest neighbor computations.

Parameters:

• x (Tensor) – First set of vectors of shape (B, D), where \(B\) is the number of samples and \(D\) is the feature dimension.

• y (Tensor) – Second set of vectors of shape (K, D), where \(K\) is the number of reference points (e.g., cluster centers).

Returns:

Squared distance matrix of shape (B, K), where element

[i, j] contains \(\| \mathbf{x}_i - \mathbf{y}_j \|_2^2\).

Return type:

Tensor

Example

>>> x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])  # 2 samples
>>> y = torch.tensor([[0.0, 0.0], [1.0, 1.0]])  # 2 centers
>>> distances = pairwise_squared_distance(x, y)
>>> print(distances)
tensor([[ 5.,  2.],
        [25., 13.]])