import torch


def sequence_mask(lengths, max_seq_length=None, dtype=torch.bool):
  """
  Returns a mask tensor representing the first N positions of each cell.

  lengths:	integer tensor, all its values <= maxlen.
  max_seq_length:	scalar integer tensor, size of last dimension of returned tensor.
    Default is the maximum value in lengths.
  dtype:	output type of the resulting tensor.
  return: A mask tensor of shape lengths.shape + (maxlen,), cast to specified dtype.

  sequence_mask([1, 3, 2], 5)  # [[True, False, False, False, False],
                                #  [True, True, True, False, False],
                                #  [True, True, False, False, False]]

  sequence_mask([[1, 3],[2,0]])  # [[[True, False, False],
                                  #   [True, True, True]],
                                  #  [[True, True, False],
                                  #   [False, False, False]]]
  """
  if max_seq_length is None:
    max_seq_length = lengths.max()
  matrix = torch.unsqueeze(lengths, dim=-1)
  row_vector = torch.arange(0, max_seq_length, 1).to(matrix.device)
  mask = row_vector < matrix

  mask.type(dtype)
  return mask