import torch


def kld_loss(mu, logvar, mean_reduction=True):
    """ First it computes the KLD over each datapoint in the batch as a sum over all latent dims. 
        It returns the mean KLD over the batch size.
        The KLD is computed in comparison to a multivariate Gaussian with zero mean and identity covariance.

    Args:
        mu (torch.Tensor): the part of the latent vector that corresponds to the mean
        logvar (torch.Tensor): the log of the variance (sigma squared)
        mean_reduction (bool): Whether to perform mean reduction across batch (default is true)

    Returns:
        (torch.Tensor): KL divergence.
    """
    mu = mu.flatten(1)
    logvar = logvar.flatten(1)
    kld_per_sample = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim = 1)
    if mean_reduction:
        return torch.mean(kld_per_sample, dim = 0)
    else:
        return kld_per_sample