import torch


def get_mask(batch_size: int, max_size: int, lengths: torch.LongTensor):
    """ Returns mask given the lengths tensor. A convenience method

    Given a lengths tensor as in

    .. code-block:: python

        >> torch.LongTensor([3, 1, 2])

    which often indicates the original length of the tensor
    without padding, `get_mask()` returns a tensor with 1 positions
    where there is no padding and 0 where there is padding

    Parameters
    ----------
    batch_size : int
        Batch size of the tensors
    max_size : int
        Maximum size or often Maximum number of time steps
    lengths : torch.LongTensor
        The original length of the tensors in the batch without padding

    Returns
    -------
    torch.LongTensor
        Mask having 1 where there are no paddings and 0 where there are paddings
    """
    assert batch_size == lengths.size(0)
    mask = []

    for length in lengths:
        zero_row = torch.zeros(max_size, dtype=torch.long)
        zero_row[: length.item()] = 1
        mask.append(zero_row.unsqueeze(0))

    mask = torch.cat(mask, dim=0)
    mask = torch.LongTensor(mask)
    return mask