code: def gen_mem_attn_mask(
      mem,
      mem_valid_length,
      data,
      data_valid_length=None,
      layout: str = "NT",
  ):
      """Generate the mask used for the decoder. All query slots are attended to the memory slots.
  
      In our implementation, 1 --> not masked, 0 --> masked
      Let's consider the data + mem with a batch of two samples:
  
      mem = [['I',   'can', 'now',   'use'],
             ['May', 'the', 'force', '<PAD>']]
      mem_valid_length =
          [4, 3]
      data =
          [['numpy', 'in',    'Gluon@@', 'NLP'  ],
           ['be',    'with',  'you',     '<PAD>']]
      data_valid_length =
          [4, 3]
  
      For our example, the mask of the first sample is
                     ['I', 'can', 'now', 'use']
          'numpy':     1,    1,     1,     1
          'in':        1,    1,     1,     1
          'Gluon@@':   1,    1,     1,     1
          'NLP':       1,    1,     1,     1
      The mask of the second sample is
                     ['be', 'with', 'you', '<PAD>']
          'May':        1,    1,     1,     0
          'the':        1,    1,     1,     0
          'force':      1,    1,     1,     0
          '<PAD>':      0,    0,     0,     0
  
      Parameters
      ----------
      mem
         - layout = 'NT'
              Shape (batch_size, mem_length, C_mem)
         - layout = 'TN'
              Shape (mem_length, batch_size, C_mem)
      mem_valid_length :
          Shape (batch_size,)
      data
          - layout = 'NT'
              Shape (batch_size, query_length, C_data)
          - layout = 'TN'
              Shape (query_length, batch_size, C_data)
      data_valid_length :
          Shape (batch_size,)
      layout
          Layout of the data + mem tensor
  
      Returns
      -------
      mask
          Shape (batch_size, query_length, mem_length)
      """
      device = mem.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)
          )
      batch_size = mem.shape[batch_axis]
      mem_length = mem.shape[time_axis]
      query_length = data[time_axis]
      mem_steps = th.arange(mem_length, device=device)  # (mem_length,)
      data_steps = th.arange(
          data.shape[time_axis], device=device
      )  # (query_length,)
      # mem_mask will have shape (B, 1, mem_length)
      mem_mask = mem_steps.view(
          (1, 1, mem_length)
      ) < mem_valid_length.view((batch_size, 1, 1))
      if data_valid_length is not None:
          # (B, query_length, 1)
          data_mask = data_steps.view(
              (1, -1, 1)
          ) < data_valid_length.view((batch_size, 1, 1))
          mask = mem_mask * data_mask
      else:
          mask = mem_mask.expand(batch_size, query_length, -1)
      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 2 for calling a global function with example usage: gen_mem_attn_mask(mem, mem_valid_length, data, data_valid_length, layout) and returns: mask