import torch


def seq_mask_from_lens(lengths, max_len=None):
    """
    Creates a boolean mask from sequence lengths.
    If the `lengths` is of shape (...), the `mask` is of shape (..., max_len).
    The last dimension is of shape (max_len) and consisting of consecutive
    `True`s and `False`s. The number of `True`s is decided by the number in
    the `lengths`. `True` means that the corresponding position is not
    padding token, and `False` otherwise.

    lengths: tensor containing the lengths of sequences
    max_len: the max length of all the sequences
    """
    if max_len is None:
        max_len = lengths.max().item()
    mask = torch.arange(0, max_len, dtype=torch.long).type_as(lengths)
    mask = mask.unsqueeze(0)
    mask = mask.repeat(*lengths.size(), 1)
    mask = mask.lt(lengths.unsqueeze(-1))
    return mask