import torch


def kl_divergence(mean, logvar, raxis=1):
    """
    KL divergence between N(mean,std) and the standard normal N(0,1).
    Args:
        mean: mean of a normal dist.
        logvar: log variance (log(std**2)) of a normal dist.
    Returns Kl divergence in batch shape.
    """
    kl_term = 1/2 * torch.sum((logvar.exp() + mean.pow(2) - logvar - 1), dim=raxis)
    
    return kl_term.unsqueeze(-1)