def gauss_kl_white_diag(q_mu, q_sqrt):
    """
    Compute the KL divergence from

          q(x) = N(q_mu, q_sqrt^2)
    to
          p(x) = N(0, I)

    We assume multiple independent distributions, given by the columns of
    q_mu and q_sqrt

    q_mu is a matrix, each column contains a mean

    q_sqrt is a matrix, each column represents the diagonal of a square-root
        matrix of the covariance.
    """

    KL = 0.5 * (q_mu**2).sum()                    # Mahalanobis term
    KL += -0.5 * q_sqrt.numel()
    KL = KL - q_sqrt.abs().log().sum()            # Log-det of q-cov
    KL += 0.5 * (q_sqrt**2).sum()                 # Trace term
    return KL