import torch


def lengths_to_mask(lengths, max_len: int, dim: int = 0, negative_mask: bool = False):
    """
    Convert a tensor of lengths to mask
    For example, lengths = [[2, 3, 4]], max_len = 5
    mask =
       [[1, 1, 1],
        [1, 1, 1],
        [0, 1, 1],
        [0, 0, 1],
        [0, 0, 0]]
    """
    assert len(lengths.size()) <= 2
    if len(lengths) == 2:
        if dim == 1:
            lengths = lengths.t()
        lengths = lengths
    else:
        lengths = lengths.unsqueeze(1)

    # lengths : batch_size, 1
    lengths = lengths.view(-1, 1)

    batch_size = lengths.size(0)
    # batch_size, max_len
    mask = (
        torch.arange(max_len, device="cuda")
        .expand(batch_size, max_len)
        .type_as(lengths)
        < lengths
    )

    if negative_mask:
        mask = ~mask

    if dim == 0:
        # max_len, batch_size
        mask = mask.t()

    return mask