Pytorch masked `log_sum_exp`

log_sum_exp.py

def log_sum_exp(inputs, keepdim=False, mask=None):
    """Numerically stable logsumexp on the last dim of `inputs`.
       reference: https://github.com/pytorch/pytorch/issues/2591
    Args:
        inputs: A Variable with any shape.
        keepdim: A boolean.
        mask: A mask variable of type float. It has the same shape as `inputs`.
              **ATTENTION** invalid entries are masked to **ONE**, not ZERO
    Returns:
        Equivalent of log(sum(exp(inputs), keepdim=keepdim)).
    """

    if mask is not None:
        mask = 1. - mask
        max_offset = -1e7 * mask
    else:
        max_offset = 0.

    s, _ = torch.max(inputs + max_offset, dim=-1, keepdim=True)

    inputs_offset = inputs - s
    if mask is not None:
        inputs_offset.masked_fill_(mask.byte(), -float('inf'))

    outputs = s + inputs_offset.exp().sum(dim=-1, keepdim=True).log()

    if not keepdim:
        outputs = outputs.squeeze(-1)
    return outputs