def mahalanobis_squared(xi, xj, VI=None):
    """Computes the pair-wise squared mahalanobis distance matrix as:

        (xi - xj)^T V^-1 (xi - xj)

    Args:
        xi (Tensor): xi input matrix.
        xj (Tensor): xj input matrix.
        VI (Tensor): The inverse of the covariance matrix, default: identity
            matrix.

    Returns:
        Weighted matrix of all pair-wise distances (Tensor).
    """
    if VI is None:
        xi_VI = xi
        xj_VI = xj
    else:
        xi_VI = xi.mm(VI)
        xj_VI = xj.mm(VI)

    D = (xi_VI * xi).sum(dim=-1).reshape(-1, 1) \
      + (xj_VI * xj).sum(dim=-1).reshape(1, -1) \
      - 2 * xi_VI.mm(xj.t())
    return D