import torch


def my_log_sum_exp(x, axis=None):
    """
    Apply log-sum-exp with subtraction of the largest element to improve numerical stability.
    """
    (x_max, idx) = torch.max(x, dim=axis, keepdim=True)

    return torch.log(torch.sum(torch.exp(x - x_max), dim=axis, keepdim=True)) + x_max