import torch


def lengths_to_mask(lengths, max_len):
    """
    This mask has a 1 for every location where we should calculate a loss

    lengths include the null terminator. for example the following is length 1:

    0 0 0 0 0 0

    The following are each length 2:

    1 0 0 0 0
    3 0 0 0 0

    if max_len is 3, then the tensor will be 3 wide. The longest tensors will look
    like:

    1 2 0
    3 4 0
    5 7 0

    Whils the rnn might not generate the final 0 each time, this is an error

    The mask for these length 3 utterances will be all 1s:

    1 1 1
    1 1 1

    """
    assert len(lengths.size()) == 1
    N = lengths.size()[0]
    cumsum = torch.zeros(N, max_len, device=lengths.device, dtype=torch.int64).fill_(1)
    cumsum = cumsum.cumsum(dim=-1)
    l_expand = lengths.view(N, 1).expand(N, max_len)
    in_alive_mask = l_expand > cumsum - 1
    return in_alive_mask