import torch


def _pairwise_distances(embeddings, squared=False):
    """Calculate pairwise distances of the given embeddings
    Args:
        embeddings (torch.tensor): Embeddings of shape (batch_size, embed_dim)
        squared (bool): Squared euclidean distance matrix
    Returns:
        dists (torch.tensor): Pairwise distances of shape (batch_size, batch_size)
    """
    dot_prod = torch.matmul(embeddings, embeddings.T)
    sq_norm = dot_prod.diagonal(0)
    dists = sq_norm.unsqueeze(0) - 2.0 * dot_prod + sq_norm.unsqueeze(1)

    # Due to computation errors some dists may be negative so we make them 0.0
    dists = torch.clamp(dists, min=0.0)

    if not squared:
        # Gradient of sqrt is infinite when dists are 0.0
        mask = dists.eq(0.0).float()
        dists = dists + mask * 1e-16
        dists = (1.0 - mask) * torch.sqrt(dists)

    return dists