import torch


def compute_inner_distmat(features: torch.Tensor):
    """
    Used to obtain `gallery_distmat` or `all_distmat`

    Args:
        features: Tensor(x, m) (x = g for gallery_distmat, x = q + g for all_distmat)

    Returns:
        Tensor(x, x)
        Euclidean squared distance
    """
    n, m = features.shape
    ff = features.pow(2).sum(dim=1, keepdim=True).expand(n, n)
    distmat = (ff + ff.t()).addmm(mat1=features, mat2=features.t(), beta=1, alpha=-2)

    return distmat