import torch


def upper_triangular_mask(seq_length):
  """
  Return a (seq_length x seq_length) mask which all upper triangular element are true,
    including the diagonal.
  :param seq_length: specify the shape of the matrix
  :return: A mask with shape (seq_length, seq_length)

  '''
  seq_length = 5
  [[ True  True  True  True  True]
   [False  True  True  True  True]
   [False False  True  True  True]
   [False False False  True  True]
   [False False False False  True]]
  '''
  """
  return torch.logical_and(
    seq_length.unsqueeze(1),  # [batch_size, 1, seq_length]
    seq_length.unsqueeze(2)  # [batch_size, seq_length, 1]
  )