import torch


def log_interpolate(log_a, log_b, alpha_logit):
    """Numerically stable implementation of log(alpha * a + (1-alpha) * b)."""
    log_alpha = -torch.nn.functional.softplus(torch.tensor(-alpha_logit))
    log_1_minus_alpha = -torch.nn.functional.softplus(torch.tensor(alpha_logit))
    y = torch.logsumexp(torch.stack((log_alpha + log_a, log_1_minus_alpha + log_b),0), dim=0)
    return y