import torch


def sequence_mask(lengths, max_len=None):
    """Creates a boolean mask from sequence lengths.
    Args:
        lengths (torch.Tensor): lengths with shape (bs,)
        max_len (int, optional): max sequence length.
            if None it will be setted to lengths.max()
    Returns:
        torch.Tensor: (bs, max_len)
    """
    if max_len is None:
        max_len = lengths.max()
    aranges = torch.arange(max_len).repeat(lengths.shape[0], 1)
    aranges = aranges.to(lengths.device)
    return aranges < lengths.unsqueeze(1)