import torch


def kl_loss(mean1, logvar1, mean2, logvar2):
    """
    KL divergence of two multivariate normal distributions with diagonal covariance.
    :param mean1: mean of distribution 1
    :param logvar1: logarithm of the covariance diagonal of distribution 1
    :param mean2: mean of distribution 2
    :param logvar2: logarithm of the covariance diagonal of distribution 2
    :return: KL divergence of distribution 1 and 2
    """
    result = -0.5 * torch.sum(logvar1 - logvar2 - torch.pow(mean1 - mean2, 2) / logvar2.exp() -
                              torch.exp(logvar1 - logvar2) + 1, 1)
    return result.mean()