import torch


def kullback_leibner(p, q, batch=False):
    """Calculates Kullback-Leibner Divergence of two probability densities

    Args:
      p, q:   density values (must be all in range [0..1])
      batch:  if True, will return summed KL, divided by batch shape

    Returns:
      KL divergence, summed and divided by batch size if `batch` is `True`
    """
    kl = p * torch.log(p / q)
    if batch:
        kl = torch.sum(kl, 0) / kl.shape[0]
    return kl