import torch


def compute_pdist_matrix(batch, p=2.0):
    """
    Computes the matrix of pairwise distances w.r.t. p-norm
    :param batch: torch.Tensor, input vectors
    :param p:     float, norm parameter, such that ||x||p = (sum_i |x_i|^p)^(1/p)
    :return: torch.Tensor, matrix A such that A_ij = ||batch[i] - batch[j]||_p (for flattened batch[i] and batch[j])
    """
    mat = torch.zeros(batch.shape[0], batch.shape[0], device=batch.device)
    ind = torch.triu_indices(batch.shape[0], batch.shape[0], offset=1, device=batch.device)
    mat[ind[0], ind[1]] = torch.pdist(batch.view(batch.shape[0], -1), p=p)

    return mat + mat.transpose(0, 1)