docstr: Parameters
  ----------
  data :
      - layout = 'NT'
          Shape (batch_size, seq_length, C_in)
      - layout = 'TN'
          Shape (seq_length, batch_size, C_in)
  mem :
      - layout = 'NT'
          Shape (batch_size, mem_length, C_mem)
      - layout = 'TN'
          Shape (mem_length, batch_size, C_mem)
  self_causal_mask :
      Shape (batch_size, seq_length, seq_length)
      Mask for the causal self-attention.
      self_causal_mask[i, j, :] masks the elements that token `j` attends to.
      To understand the self-causal attention mask, we can look at the following example:
                 ['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
  mem_attn_mask :
      Shape (batch_size, seq_length, mem_length)
      Mask between the decoding input and the memory.
                 ['numpy', 'in', 'Gluon@@', 'NLP']
      'I':         1,     1,      1,      1
      'can':       1,     1,      1,      1
      'now':       1,     1,      1,      1
      'use':       1,     1,      1,      1
  
  Returns
  -------
  out :
      - layout = 'NT'
          Shape (batch_size, seq_length, C_out)
      - layout = 'TN'
          Shape (seq_length, batch_size, C_out)
docstr_tok: ["Parameters","--","--","--","--","--","data",":","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",",","C_in",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",",","C_in",")","mem",":","-","layout","=","'NT'","Shape","(","batch_size",",","mem_length",",","C_mem",")","-","layout","=","'TN'","Shape","(","mem_length",",","batch_size",",","C_mem",")","self_causal_mask",":","Shape","(","batch_size",",","seq_length",",","seq_length",")","Mask","for","the","causal","self-attention",".","self_causal_mask","[","i",",","j",",",":","]","masks","the","elements","that","token","`","j","`","attends","to",".","To","understand","the","self-causal","attention","mask",",","we","can","look","at","the","following","example",":","[","'","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","mem_attn_mask",":","Shape","(","batch_size",",","seq_length",",","mem_length",")","Mask","between","the","decoding","input","and","the","memory",".","[","'numpy","'",",","'in","'",",","'Gluon","@","@","'",",","'NLP","'","]","'","I","'",":","1",",","1",",","1",",","1","'can","'",":","1",",","1",",","1",",","1","'now","'",":","1",",","1",",","1",",","1","'use","'",":","1",",","1",",","1",",","1","Returns","--","--","--","-","out",":","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",",","C_out",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",",","C_out",")"]
code: def forward(self, data, mem, self_causal_mask, mem_attn_mask):
      """
      Parameters
      ----------
      data :
          - layout = 'NT'
              Shape (batch_size, seq_length, C_in)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_in)
      mem :
          - layout = 'NT'
              Shape (batch_size, mem_length, C_mem)
          - layout = 'TN'
              Shape (mem_length, batch_size, C_mem)
      self_causal_mask :
          Shape (batch_size, seq_length, seq_length)
          Mask for the causal self-attention.
          self_causal_mask[i, j, :] masks the elements that token `j` attends to.
          To understand the self-causal attention mask, we can look at the following example:
                     ['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
      mem_attn_mask :
          Shape (batch_size, seq_length, mem_length)
          Mask between the decoding input and the memory.
                     ['numpy', 'in', 'Gluon@@', 'NLP']
          'I':         1,     1,      1,      1
          'can':       1,     1,      1,      1
          'now':       1,     1,      1,      1
          'use':       1,     1,      1,      1
  
      Returns
      -------
      out :
          - layout = 'NT'
              Shape (batch_size, seq_length, C_out)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_out)
      """
      # 1. Get the causal self-attention value
      residual = data
      if self._pre_norm:
          data = self.ln_in(data)
      self_query, self_key, self_value = th.split(
          self.attn_in_qkv(data), self._units, dim=-1
      )
      out, [_, self_attn_weight] = self.self_attention(
          self_query.reshape(
              (
                  self_query.shape[0],
                  self_query.shape[1],
                  self._num_heads,
                  -1,
              )
          ),
          self_key.reshape(
              (
                  self_key.shape[0],
                  self_key.shape[1],
                  self._num_heads,
                  -1,
              )
          ),
          self_value.reshape(
              (
                  self_value.shape[0],
                  self_value.shape[1],
                  self._num_heads,
                  -1,
              )
          ),
          self_causal_mask,
      )
      out = self.proj_in(out)
      out = self.dropout_layer(out)
      out = out + residual
      if not self._pre_norm:
          out = self.ln_in(out)
      # 2. Attend to the contextual memory
      data = out
      residual = data
      if self._pre_norm:
          data = self.ln_inter(data)
      out, [_, context_attn_weight] = self.inter_attention(
          th.reshape(
              self.attn_inter_q(data),
              (data.shape[0], data.shape[1], self._num_heads, -1),
          ),
          th.reshape(
              self.attn_inter_k(mem),
              (mem.shape[0], mem.shape[1], self._num_heads, -1),
          ),
          th.reshape(
              self.attn_inter_v(mem),
              (mem.shape[0], mem.shape[1], self._num_heads, -1),
          ),
          mem_attn_mask,
      )
      out = self.proj_inter(out)
      out = self.dropout_layer(out)
      out = out + residual
      if not self._pre_norm:
          out = self.ln_inter(out)
      # 3. Encode the output via an FFN layer
      out = self.ffn(out)
      return out
  
imports: from typing import Union, Optional, List
  from torch import nn
  import torch
  from ..layers import PositionwiseFFN
  from ..attention_cell import MultiHeadAttentionCell, gen_mem_attn_mask, gen_self_attn_mask
  from ..utils import to_torch_dtype
label_desc: Use image node_id 5 for calling the TransformerDecoderLayer obj's underlying member method code with example usage: obj.forward(data, mem, _causal_mask, mem_attn_mask) and returns: out