import torch


def KLD_cost(mu_p, sig_p, mu_q, sig_q):
    """KLD between two gaussian distribution

    Parameters
    ----------
    mu_p : :obj:`torch.tensor`
    sig_p : :obj:`torch.tensor`
    mu_q : :obj:`torch.tensor`
    sig_q : :obj:`torch.tensor`

    Returns
    -------
    :obj:`torch.tensor`
        KL divergence
    """
    KLD = 0.5 * (2 * torch.log(sig_p / sig_q) - 1 + (sig_q / sig_p).pow(2)
                 + ((mu_p - mu_q) / sig_p).pow(2)).sum()
    # https://arxiv.org/abs/1312.6114 
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    return KLD