import torch


def softmax(X):
    """
    Entropy-smoothed max, a.k.a. logsumexp.

    Solves $max_{p \in \Delta^d} <x, p> - \sum_{i=1}^d p_i \log(p_i)$ along
    dim=2.

    :param x: torch.Tensor, shape = (b, n, m)
        Vector to project

    :return: torch.Tensor, shape = (b, n)
        Projected vector
    """
    M, _ = torch.max(X, dim=2)
    X = X - M[:, :, None]
    S = torch.sum(torch.exp(X), dim=2)
    M = M + torch.log(S)
    return M