import torch


def kl_gaussian_mog(p, q, num_samples=1):
    """
    Estimate KL(p||q) = E_p[ \log q(z) - log p(z)] = E_p[\log p(z)] - E_p[log q(z)] = - H(p) - E_p[log q(z)]
        where we either use the closed-form entropy of p if available, or MC estimate it, 
        and always MC-estimate the second term.
    """
    if num_samples == 1:
        # [1, ...]
        z = p.rsample().unsqueeze(0)
    else:
        # [num_samples, ...]
        z = p.rsample(torch.Size([num_samples]))
    try:
        H_p = p.entropy()
    except NotImplementedError:
        H_p = - p.log_prob(z).mean(0)

    return - H_p - q.log_prob(z).mean(0)