import torch


def compute_logits(cluster_centers, data):
    """Computes the logits of being in one cluster, squared Euclidean.
    Args:
    cluster_centers: [K, D] Cluster center representation.
    data: [B, N, D] Data representation.
    Returns:
    log_prob: [B, N, K] logits.
    """
    k = cluster_centers.shape[0]
    b = data.shape[0]
    cluster_centers = cluster_centers.unsqueeze(dim=0)  # [1, K, D]
    data = data.contiguous().view(-1, data.shape[-1]).unsqueeze(dim=1)  # [N, 1, D]
    # neg_dist = -torch.sum(torch.pow(data - cluster_centers, 2), dim=-1)  # [N, K]
    neg_dist = -torch.mean(torch.pow(data - cluster_centers, 2), dim=-1)  # [N, K]
    neg_dist = neg_dist.view(b, -1, k)
    return neg_dist