def bernoulli_kl(q_probs, p_probs):
    """
        computes the KL divergence per batch between two Bernoulli distributions
        parametrized by q_probs and p_probs

    Note: EPS is added for numerical stability, see
    https://github.com/pytorch/pytorch/issues/15288

    Args:
        q_probs (torch tensor): mean of q [batch_size, *latent_dims]
        p_probs (torch tensor): mean of p [batch_size, *latent_dims]

    Returns:
        kl_div (torch tensor): kl divergence [batch_size, *latent_dims]
    """
    EPS = 1e-32
    p1 = p_probs
    p0 = 1 - p1
    q1 = q_probs
    q0 = 1 - q1

    logq1 = (q1 + EPS).log()
    logq0 = (q0 + EPS).log()
    logp1 = (p1).log()
    logp0 = (p0).log()

    kl_div_1 = q1 * (logq1 - logp1)
    kl_div_0 = q0 * (logq0 - logp0)
    return kl_div_1 + kl_div_0