docstr: Parameters
  ----------
  data
      The input data
  
      - layout = 'NT'
          Shape (B, T)
      - layout = 'TN'
          Shape (T, B)
  
  target
      The ground truth
  
      - layout = 'NT'
          Shape (B, T)
      - layout = 'TN'
          Shape (T, B)
  
  mem_l
      A list of memory objects
  
      - layout = 'NT'
          Shape (B, T_mem, units)
      - layout = 'TN'
          Shape (T_mem, B, units)
  
  rel_positions
      Shape (query_length, mem_length + query_length)
      By default, we will use the following relative positions
  
      .. code-block:: none
  
                     ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
          'in':        5,    4,     3,     2,      1,     0,      -1,      -2
          'Gluon@@':   6,    5,     4,     3,      2,     1,       0,      -1
          'NLP':       7,    6,     5,     4,      3,     2,       1,       0
  
  data_mem_mask
      Shape (B, query_length, mem_length + query_length)
      Here, 1 --> will be used, 0 --> won't be used.
      By default, we will mask all locations that have distance > mem_length with the
      current token.
      Following is an example in which query_length = 3, mem_length = 4
  
      .. code-block:: none
  
                      |------- <mem> ----------|--------- <query> ------------|
           <query>   ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
          'numpy':     1,    1,     1,     1,      1,     0,      0,        0
          'in':        0,    1,     1,     1,      1,     1,      0,        0
          'Gluon@@':   0,    0,     1,     1,      1,     1,      1,        0
          'NLP':       0,    0,     0,     1,      1,     1,      1,        1
  
      Also, we provide the option in which we only mask the future tokens, this is
      supported by setting `causal_only` to True. However, there will be a
      discrepancy between training and inference because the effecitve memory length is
      longer for the later tokens in the query.
  
      .. code-block:: none
  
                      |------- <mem> ----------|--------- <query> ------------|
           <query>   ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
          '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
  
  causal_only
      Whether to ignore the local masking constraint. See the flag above for more information.
  
  detach_memory
      Whether to detach the encoded memory from the graph.
  
  Returns
  -------
  logits
      The selected logits
  
      - layout = 'NT'
          Shape (B, T)
      - layout = 'TN'
          Shape (T, B)
  
  new_mem_l
      A list of the updated memory
  
      - layout = 'NT'
          Each will have shape (B, T, C)
      - layout = 'TN'
          Each will have shape (T, B, C)
docstr_tok: ["Parameters","--","--","--","--","--","data","The","input","data","-","layout","=","'NT'","Shape","(","B",",","T",")","-","layout","=","'TN'","Shape","(","T",",","B",")","target","The","ground","truth","-","layout","=","'NT'","Shape","(","B",",","T",")","-","layout","=","'TN'","Shape","(","T",",","B",")","mem_l","A","list","of","memory","objects","-","layout","=","'NT'","Shape","(","B",",","T_mem",",","units",")","-","layout","=","'TN'","Shape","(","T_mem",",","B",",","units",")","rel_positions","Shape","(","query_length",",","mem_length","+","query_length",")","By","default",",","we","will","use","the","following","relative","positions","..","code-block",":",":","none","[","'","I","'",",","'can","'",",","'now","'",",","'use","'",",","'numpy","'",",","'in","'",",","'Gluon","@","@","'",",","'NLP","'","]","'in","'",":","5",",","4",",","3",",","2",",","1",",","0",",","-1",",","-2","'Gluon","@","@","'",":","6",",","5",",","4",",","3",",","2",",","1",",","0",",","-1","'NLP","'",":","7",",","6",",","5",",","4",",","3",",","2",",","1",",","0","data_mem_mask","Shape","(","B",",","query_length",",","mem_length","+","query_length",")","Here",",","1","--",">","will","be","used",",","0","--",">","wo","n't","be","used",".","By","default",",","we","will","mask","all","locations","that","have","distance",">","mem_length","with","the","current","token",".","Following","is","an","example","in","which","query_length","=","3",",","mem_length","=","4","..","code-block",":",":","none","|","--","--","--","-","<","mem",">","--","--","--","--","--","|","--","--","--","--","-","<","query",">","--","--","--","--","--","--","|","<","query",">","[","'","I","'",",","'can","'",",","'now","'",",","'use","'",",","'numpy","'",",","'in","'",",","'Gluon","@","@","'",",","'NLP","'","]","'numpy","'",":","1",",","1",",","1",",","1",",","1",",","0",",","0",",","0","'in","'",":","0",",","1",",","1",",","1",",","1",",","1",",","0",",","0","'Gluon","@","@","'",":","0",",","0",",","1",",","1",",","1",",","1",",","1",",","0","'NLP","'",":","0",",","0",",","0",",","1",",","1",",","1",",","1",",","1","Also",",","we","provide","the","option","in","which","we","only","mask","the","future","tokens",",","this","is","supported","by","setting","`","causal_only","`","to","True",".","However",",","there","will","be","a","discrepancy","between","training","and","inference","because","the","effecitve","memory","length","is","longer","for","the","later","tokens","in","the","query",".","..","code-block",":",":","none","|","--","--","--","-","<","mem",">","--","--","--","--","--","|","--","--","--","--","-","<","query",">","--","--","--","--","--","--","|","<","query",">","[","'","I","'",",","'can","'",",","'now","'",",","'use","'",",","'numpy","'",",","'in","'",",","'Gluon","@","@","'",",","'NLP","'","]","'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","causal_only","Whether","to","ignore","the","local","masking","constraint",".","See","the","flag","above","for","more","information",".","detach_memory","Whether","to","detach","the","encoded","memory","from","the","graph",".","Returns","--","--","--","-","logits","The","selected","logits","-","layout","=","'NT'","Shape","(","B",",","T",")","-","layout","=","'TN'","Shape","(","T",",","B",")","new_mem_l","A","list","of","the","updated","memory","-","layout","=","'NT'","Each","will","have","shape","(","B",",","T",",","C",")","-","layout","=","'TN'","Each","will","have","shape","(","T",",","B",",","C",")"]
code: def forward(
      self,
      data,
      target,
      mem_l,
      rel_positions=None,
      data_mem_mask=None,
      causal_only=False,
      detach_memory=True,
  ):
      """
  
      Parameters
      ----------
      data
          The input data
  
          - layout = 'NT'
              Shape (B, T)
          - layout = 'TN'
              Shape (T, B)
  
      target
          The ground truth
  
          - layout = 'NT'
              Shape (B, T)
          - layout = 'TN'
              Shape (T, B)
  
      mem_l
          A list of memory objects
  
          - layout = 'NT'
              Shape (B, T_mem, units)
          - layout = 'TN'
              Shape (T_mem, B, units)
  
      rel_positions
          Shape (query_length, mem_length + query_length)
          By default, we will use the following relative positions
  
          .. code-block:: none
  
                         ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
              'in':        5,    4,     3,     2,      1,     0,      -1,      -2
              'Gluon@@':   6,    5,     4,     3,      2,     1,       0,      -1
              'NLP':       7,    6,     5,     4,      3,     2,       1,       0
  
      data_mem_mask
          Shape (B, query_length, mem_length + query_length)
          Here, 1 --> will be used, 0 --> won't be used.
          By default, we will mask all locations that have distance > mem_length with the
          current token.
          Following is an example in which query_length = 3, mem_length = 4
  
          .. code-block:: none
  
                          |------- <mem> ----------|--------- <query> ------------|
               <query>   ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
              'numpy':     1,    1,     1,     1,      1,     0,      0,        0
              'in':        0,    1,     1,     1,      1,     1,      0,        0
              'Gluon@@':   0,    0,     1,     1,      1,     1,      1,        0
              'NLP':       0,    0,     0,     1,      1,     1,      1,        1
  
          Also, we provide the option in which we only mask the future tokens, this is
          supported by setting `causal_only` to True. However, there will be a
          discrepancy between training and inference because the effecitve memory length is
          longer for the later tokens in the query.
  
          .. code-block:: none
  
                          |------- <mem> ----------|--------- <query> ------------|
               <query>   ['I', 'can', 'now', 'use', 'numpy', 'in', 'Gluon@@', 'NLP']
              '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
  
      causal_only
          Whether to ignore the local masking constraint. See the flag above for more information.
  
      detach_memory
          Whether to detach the encoded memory from the graph.
  
      Returns
      -------
      logits
          The selected logits
  
          - layout = 'NT'
              Shape (B, T)
          - layout = 'TN'
              Shape (T, B)
  
      new_mem_l
          A list of the updated memory
  
          - layout = 'NT'
              Each will have shape (B, T, C)
          - layout = 'TN'
              Each will have shape (T, B, C)
      """
      # Note that curr_mem_length will not necessarily be equal to mem_length
      if self._layout == "NT":
          time_axis = 1
          batch_axis = 0
      elif self._layout == "TN":
          time_axis = 0
          batch_axis = 1
      else:
          raise NotImplementedError
      query_length = data.shape[time_axis]
      curr_mem_length = mem_l[0].shape[time_axis]
      batch_size = mem_l[0].shape[batch_axis]
      ctx = data.ctx
      local_attn_mask = mx.np.ones(
          (batch_size, query_length, curr_mem_length + query_length),
          dtype=np.int32,
          ctx=ctx,
      )
      if not causal_only:
          # Generate the mask, we mask out the input outside the local self.mem_length window
          local_attn_mask = mx.np.triu(
              mx.np.tril(local_attn_mask, curr_mem_length),
              curr_mem_length - self.mem_length,
          )
      else:
          local_attn_mask = mx.np.tril(local_attn_mask, curr_mem_length)
      if data_mem_mask is None:
          data_mem_mask = local_attn_mask
      else:
          data_mem_mask = data_mem_mask * local_attn_mask
      if rel_positions is None:
          query_ids = mx.np.arange(
              curr_mem_length,
              curr_mem_length + query_length,
              dtype=np.int32,
              ctx=ctx,
          )
          mem_ids = mx.np.arange(
              0, curr_mem_length + query_length, dtype=np.int32, ctx=ctx
          )
          rel_positions = mx.np.expand_dims(
              query_ids, axis=1
          ) - mx.np.expand_dims(mem_ids, axis=0)
      # Get word embeddings
      word_embeddings = self.word_emb(data)
      word_embeddings = self.dropout_layer(word_embeddings)
      out_l = self.decoder(
          word_embeddings, mem_l, rel_positions, data_mem_mask
      )
      # Get the output logits
      logits = self.crit(out_l[-1], target)
  
      # Get the new memory
      new_mem_l = []
      for step_out, mem in zip([word_embeddings] + out_l, mem_l):
          new_mem = mx.np.concatenate([mem, step_out], axis=time_axis)
          if self._layout == "NT":
              new_mem = new_mem[:, -self.mem_length :]
          elif self._layout == "TN":
              new_mem = new_mem[-self.mem_length :, :]
          else:
              raise NotImplementedError
          if detach_memory:
              new_mem_l.append(new_mem.detach())
          else:
              new_mem_l.append(new_mem)
      return logits, new_mem_l
  
imports: import numpy
  import mxnet
  from mxnet import use_np, np, npx
  from mxnet.gluon import nn, Block, HybridBlock, Parameter
  from ..attention_cell import multi_head_dot_attn, gen_self_attn_mask, gen_mem_attn_mask, RelAttentionScoreCell, MultiHeadAttentionCell
  from ..layers import get_activation, PositionalEmbedding, PositionwiseFFN, AdaptiveEmbedding, ProjectedAdaptiveLogSoftmaxWithLoss
  from ..utils.config import CfgNode
  from ..sequence_sampler import BaseStepDecoder
label_desc: Use image node_id 9 for calling the TransformerXLForLM obj's underlying member method code with example usage: obj.forward(data, target, mem_l, rel_positions, data_mem_mask, causal_only, detach_memory) and returns: logits,  new_mem_l