import torch


def gaussian_kld(mu_left, sigma_left, mu_right, sigma_right):
    """
    Provided by @anirudh
    Compute KL divergence between a bunch of univariate Gaussian distributions
    with the given means and log-variances.
    We do KL(N(mu_left, logvar_left) || N(mu_right, logvar_right)).
    """
    logsigma_left = sigma_left.log()
    logsigma_right = sigma_right.log()
    logvar_left = 2 * logsigma_left
    logvar_right = 2 * logsigma_right

    gauss_klds = 0.5 * (logvar_right - logvar_left +
                        (torch.exp(logvar_left) / torch.exp(logvar_right)) +
                        ((mu_left - mu_right) ** 2.0 / torch.exp(logvar_right)) - 1.0)
    assert len(gauss_klds.size()) == 2
    return gauss_klds
    return torch.sum(gauss_klds, 1)