code: def gen_self_attn_mask(
      data,
      valid_length=None,
      attn_type: str = "full",
      layout: str = "NT",
  ):
      """Generate the mask used for the encoder, i.e, self-attention.
  
      In our implementation, 1 --> not masked, 0 --> masked
      Let's consider the data with two samples:
      data =
          [['I',   'can', 'now',   'use', 'numpy', 'in',  'Gluon@@', 'NLP'  ],
           ['May', 'the', 'force', 'be',  'with',  'you', '<PAD>',   '<PAD>']]
      valid_length =
          [8, 6]
      - attn_type = 'causal'
          Each token will attend to itself + the tokens before.
          It will not attend to tokens in the future.
          For our example, the mask of the first sample is
                     ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
          'I':         1,    0,     0,     0,      0,     0,      0,      0
          'can':       1,    1,     0,     0,      0,     0,      0,      0
          'now':       1,    1,     1,     0,      0,     0,      0,      0
          'use':       1,    1,     1,     1,      0,     0,      0,      0
          'numpy':     1,    1,     1,     1,      1,     0,      0,      0
          'in':        1,    1,     1,     1,      1,     1,      0,      0
          'Gluon@@':   1,    1,     1,     1,      1,     1,      1,      0
          'NLP':       1,    1,     1,     1,      1,     1,      1,      1
          The mask of the second sample is
                     ['May', 'the', 'force', 'be', 'with', 'you', '<PAD>', '<PAD>']
          'May':        1,    0,     0,     0,      0,     0,      0,      0
          'the':        1,    1,     0,     0,      0,     0,      0,      0
          'force':      1,    1,     1,     0,      0,     0,      0,      0
          'be':         1,    1,     1,     1,      0,     0,      0,      0
          'with':       1,    1,     1,     1,      1,     0,      0,      0
          'you':        1,    1,     1,     1,      1,     1,      0,      0
          '<PAD>':      0,    0,     0,     0,      0,     0,      0,      0
          '<PAD>':      0,    0,     0,     0,      0,     0,      0,      0
      - attn_type = 'full'
          Each token will attend to both the tokens before and in the future
          For our example, the mask of the first sample is
                     ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
          'I':         1,    1,     1,     1,      1,     1,      1,      1
          'can':       1,    1,     1,     1,      1,     1,      1,      1
          'now':       1,    1,     1,     1,      1,     1,      1,      1
          'use':       1,    1,     1,     1,      1,     1,      1,      1
          'numpy':     1,    1,     1,     1,      1,     1,      1,      1
          'in':        1,    1,     1,     1,      1,     1,      1,      1
          'Gluon@@':   1,    1,     1,     1,      1,     1,      1,      1
          'NLP':       1,    1,     1,     1,      1,     1,      1,      1
          The mask of the second sample is
                     ['May', 'the', 'force', 'be', 'with', 'you', '<PAD>', '<PAD>']
          'May':        1,    1,     1,     1,      1,     1,      0,      0
          'the':        1,    1,     1,     1,      1,     1,      0,      0
          'force':      1,    1,     1,     1,      1,     1,      0,      0
          'be':         1,    1,     1,     1,      1,     1,      0,      0
          'with':       1,    1,     1,     1,      1,     1,      0,      0
          'you':        1,    1,     1,     1,      1,     1,      0,      0
          '<PAD>':      0,    0,     0,     0,      0,     0,      0,      0
          '<PAD>':      0,    0,     0,     0,      0,     0,      0,      0
  
      Parameters
      ----------
      data
          The data.
          - layout = 'NT'
              Shape (batch_size, seq_length, C)
          - layout = 'TN'
              Shape (seq_length, batch_size, C)
      valid_length
          Shape (batch_size,)
      attn_type
          Can be 'full' or 'causal'
      layout
          The layout of the data
  
      Returns
      -------
      mask
          Shape (batch_size, seq_length, seq_length)
      """
      device = data.device
      if layout == "NT":
          batch_axis, time_axis = 0, 1
      elif layout == "TN":
          batch_axis, time_axis = 1, 0
      else:
          raise NotImplementedError(
              "Unsupported layout={}".format(layout)
          )
      if attn_type == "full":
          if valid_length is not None:
              steps = th.arange(
                  data.shape[time_axis], device=device
              )  # (seq_length,)
              mask1 = steps.view((1, 1, -1)) < valid_length.view(
                  (valid_length.shape[0], 1, 1)
              )
              mask2 = steps.view((1, -1, 1)) < valid_length.view(
                  (valid_length.shape[0], 1, 1)
              )
              mask = mask1 * mask2
          else:
              seq_len_ones = th.ones(
                  (data.shape[time_axis],), device=device
              )  # (seq_length,)
              batch_ones = th.ones(
                  (data.shape[batch_axis],), device=device
              )  # (batch_size,)
              mask = (
                  batch_ones.view((-1, 1, 1))
                  * seq_len_ones.view((1, -1, 1))
                  * seq_len_ones.view((1, 1, -1))
              )
      elif attn_type == "causal":
          steps = th.arange(data.shape[time_axis], device=device)
          # mask: (seq_length, seq_length)
          # batch_mask: (batch_size, seq_length)
          mask = th.unsqueeze(steps, dim=0) <= th.unsqueeze(
              steps, dim=1
          )
          if valid_length is not None:
              batch_mask = th.unsqueeze(steps, dim=0) < th.unsqueeze(
                  valid_length, dim=-1
              )
              mask = mask * th.unsqueeze(batch_mask, dim=-1)
          else:
              batch_ones = th.ones(
                  data.shape[batch_axis], device=device
              )
              mask = mask * batch_ones.view((-1, 1, 1))
      else:
          raise NotImplementedError
      return mask.type(th.bool)
  
file_path: gluon-nlp/src/gluonnlp/torch/attention_cell.py
imports: import math
  import torch
  from torch import nn
  import torch.nn.functional
  from ..base import use_einsum_optimization
  from typing import Optional
label_desc: Use image node_id 1 for calling a global function with example usage: gen_self_attn_mask(data, valid_length, attn_type, layout) and returns: mask