import torch


def _pairwise_distances(vecs, squared=False):
    """
    computes matrix A where A_ij = distance(vec i, vec j)
    vecs = (N x D) => returns (N x N) matrix
    """
    norms = torch.einsum('ij,ij->i', vecs, vecs)[:, None]
    distances = norms + norms.t() - 2 * (vecs @ vecs.t())
    if not squared:
        distances = torch.sqrt(distances)
    return distances