import torch


def paired_euclidean_distance(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
    """
        Compute the distance in the semantic alignment loss (3) by 
        computing average pairwise distances between *already paired*
        points in the embedding space.

        Note this is NOT computed between all possible pairs. Rather, we
        compare i'th vector of X vs. i'th vector of Y (i == j always).

        Args:
        -   X: Pytorch tensor of shape (N,D) representing N embeddings of dim D
        -   Y: Pytorch tensor of shape (N,D) representing N embeddings of dim D

        Returns:
        -   dists: Pytorch tensor of shape (N,) representing distances between 
                fixed pairs
    """
    device = X.device
    N, D = X.shape
    assert Y.shape == X.shape
    eps = 1e-08 * torch.ones((N,1))
    eps = eps.to(device) # make sure in same memory (CPU or CUDA)
    # compare i'th vector of x vs. i'th vector of y (i == j always)
    diff = torch.pow(X - Y, 2)

    affinities = torch.sum(diff, dim=1, keepdim=True)
    # clamp the affinities to be > 1e-8 ?? Unclear why the authors do this...
    affinities = torch.max(affinities, eps)
    return torch.sqrt(affinities)