import torch


def mahalanobis_distance(x, logvar):
    """
    Returns a matrix where each element is the Mahalanobis distance between each pair of rows in x,
    assuming diagonal covariance, with covariance of each dimension shared among all examples. 
    D_{ij} = (x_i - x_j)^T Σ^{-1} (x_i - x_j)

    Parameters
    ----------
    x:      torch.Tensor [batch_size, D]
    logvar:  torch.Tensor [batch_size, D] - Diagonal of log-variance matrix
    """
    x = x.double()
    logvar = logvar.double()

    sqrt_Sigma_inv_diag = torch.exp(-0.5 * logvar)
    x = x * sqrt_Sigma_inv_diag
    xx = (x*x).sum(1, keepdim=True)
    gramian = x @ x.T

    m_dist = xx - 2.0 * gramian + xx.T

    return torch.clamp(m_dist, min=0.)