import torch


def kl_gumbel(logits, num_vertices):
    """Computes the analytical kl(q(z|x)||p(z)) = u + exp(-u) - 1.
    q(z|x) is gumbel distributed with location (u) given by logits.
    p(z) is gumbel distributed with location zero.
    """
    kl_div = logits + torch.exp(-logits) - 1.0
    return kl_div.sum((1, 2)) / num_vertices