import torch
import math


def mean_logits(logits, dim=0):
    """Return the mean logit, where the average is taken over across `dim` in probability space."""
    # p = e^logit / (sum e^logit) <=> log(p) = logit - log_sum_exp(logit)
    log_prob = logits - torch.logsumexp(logits, -1, keepdim=True)
    # mean(p) = 1/Z sum exp(log(p)) <=> log(mean(p)) = log_sum_exp(log(p)) - log(Z)
    log_mean_prob = torch.logsumexp(log_prob, 0) - math.log(log_prob.size(0))
    return log_mean_prob