import torch


def log_sum_exp(batch_logit):
    """
    Caculate the log-sum-exp operation for the last dimension.
    :param batch_logit: Size([*, logit_size]), * should at least be 1
    :return: Size([*])
    """
    batch_max, _ = batch_logit.max(dim=-1)
    batch_broadcast = batch_max.unsqueeze(-1)
    return batch_max + torch.log(
        torch.sum(torch.exp(batch_logit - batch_broadcast), dim=-1)
    )