import torch


def log_sum_exp(inp, m_size):
    """
    LogSumExp trick: https://en.wikipedia.org/wiki/LogSumExp
    Args:
        inp: size=(batch_size, vanishing_dim, hidden_dim)
        m_size: hidden_dim
    Returns:
        out: size=(batch_size, hidden_dim)
    """
    _, idx = torch.max(inp, 1)  # batch_size * hidden_dim
    max_score = torch.gather(inp, 1, idx.view(-1, 1, m_size)).view(
        -1, 1, m_size
    )  # batch_size * 1 * hidden_dim
    return max_score.view(-1, m_size) + torch.log(
        torch.sum(torch.exp(inp - max_score.expand_as(inp)), 1)
    ).view(-1, m_size)