import math
import torch


def logit_uniform_mixture(values, delta, dim=-1):
    """
    Given logits for a categorical distribution D, return the logits
    for a mixture distribution (1-delta)*D + delta*uniform
    """
    n = values.shape[dim]
    term1 = values + math.log(1 - delta)
    term2 = torch.logsumexp(values + math.log(delta / n), dim=dim, keepdim=True).expand_as(term1)
    logits = torch.stack([term1, term2]).logsumexp(dim=0)
    return logits