import torch


def sparse_categorical_kl(log_q, p_support, log_p):
    """
    Computes the restricted Kl divergence::

        sum_i restrict(q)(i) (log q(i) - log p(i))

    where ``p`` is a uniform prior, ``q`` is the posterior, and
    ``restrict(q))`` is the posterior restricted to the support of ``p`` and
    renormalized. Note for degenerate ``p=delta(i)`` this reduces to the log
    likelihood ``log q(i)``.
    """
    assert log_q.dim() == 1
    assert log_p.dim() == 1
    assert p_support.shape == log_p.shape + log_q.shape
    q = log_q.exp()
    sum_q = torch.mv(p_support, q)
    sum_q_log_q = torch.mv(p_support, q * log_q)
    sum_r_log_q = sum_q_log_q / sum_q  # restrict and normalize
    kl = sum_r_log_q - log_p  # note sum_r_log_p = log_p because p is uniform
    return kl.sum()