import torch


def logsumexp(x, dim):
    """Compute the log-sum-exp in a numerically stable way.
    Arguments
    ---------
    x : :class:`torch:torch.Tensor`
    dim : int
        The dimension along wich the operation should be computed.
    Returns
    --------
    :class:`torch:torch.Tensor`
        The dimension along which the sum is done is not squeezed.
    """
    x_max = torch.max(x, dim, keepdim=True)[0]
    return (
        torch.log(torch.sum(torch.exp(x - x_max.expand_as(x)), dim, keepdim=True))
        + x_max
    )