import torch


def gaus_kl(q_mu, q_logsigmasq, p_mu, p_logsigmasq, dim=1):
    """
    Compute KL-divergence KL(q || p) between n pairs of Gaussians
    with diagonal covariational matrices.
    Do not divide KL-divergence by the dimensionality of the latent space.

    Input: q_mu, p_mu, Tensor of shape n x d - mean vectors for n Gaussians.
    Input: q_sigma, p_sigma, Tensor of shape n x d - standard deviation
           vectors for n Gaussians.
    Return: Tensor of shape n - each component is KL-divergence between
            a corresponding pair of Gaussians.
    """
    res = p_logsigmasq - q_logsigmasq - 1 + torch.exp(q_logsigmasq - p_logsigmasq)
    res = res + (q_mu - p_mu).pow(2) * torch.exp(-p_logsigmasq)
    if dim is not None:
        return 0.5 * res.sum(dim=dim)
    else:
        return 0.5 * res