import torch


def _kl_divergence(probs1, probs2):
    """
    Computes the KL divergence in the last dimension of `probs1` and `probs2`
    as KL(P1 || P2). `probs1` and `probs2` must be the same shape. For example,
    if they are both A x B x L arrays, then the KL divergence of corresponding
    L-arrays will be computed and returned in an A x B array. Does not
    renormalize the arrays. If probs2[i] is 0, that value contributes 0.
    """

    idxs = ((probs1 != 0) & (probs2 != 0))
    quot_ = torch.divide(probs1, probs2)

    quot = torch.ones_like(probs1)
    quot[idxs] = quot_[idxs]
    return torch.sum(probs1 * torch.log(quot), dim=-1)