docstr: Parameters
  ----------
  data
      - layout = 'NT'
          Shape (batch_size, seq_length, C_in)
      - layout = 'TN'
          Shape (seq_length, batch_size, C_in)
  
  valid_length
      Shape (batch_size,)
  mem_data
      - layout = 'NT'
          Shape (batch_size, mem_length, C_mem)
      - layout = 'TN'
          Shape (mem_length, batch_size, C_mem)
  
  mem_valid_length
      Shape (batch_size,)
  
  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",")","valid_length","Shape","(","batch_size",",",")","mem_data","-","layout","=","'NT'","Shape","(","batch_size",",","mem_length",",","C_mem",")","-","layout","=","'TN'","Shape","(","mem_length",",","batch_size",",","C_mem",")","mem_valid_length","Shape","(","batch_size",",",")","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, valid_length, mem_data, mem_valid_length):
      """
  
      Parameters
      ----------
      data
          - layout = 'NT'
              Shape (batch_size, seq_length, C_in)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_in)
  
      valid_length
          Shape (batch_size,)
      mem_data
          - layout = 'NT'
              Shape (batch_size, mem_length, C_mem)
          - layout = 'TN'
              Shape (mem_length, batch_size, C_mem)
  
      mem_valid_length
          Shape (batch_size,)
  
      Returns
      -------
      out
          - layout = 'NT'
              Shape (batch_size, seq_length, C_out)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_out)
  
      """
      # 1. Embed the data
      out = self.dropout_layer(data)
      if self._data_norm:
          out = self.ln_data(out)
      self_causal_mask = gen_self_attn_mask(
          data,
          valid_length,
          dtype=self._dtype,
          attn_type="causal",
          layout=self._layout,
      )
      mem_attn_mask = gen_mem_attn_mask(
          mem_data,
          mem_valid_length,
          data,
          valid_length,
          dtype=self._dtype,
          layout=self._layout,
      )
      for i in range(self.num_layers):
          if self.recurrent:
              layer = self.layers[0]
          else:
              layer = self.layers[i]
          out = layer(out, mem_data, self_causal_mask, mem_attn_mask)
      if self._pre_norm:
          out = self.ln_final(out)
      return out
  
imports: import numpy
  import mxnet
  from mxnet import np, npx
  from mxnet import use_np
  from mxnet.gluon import nn, HybridBlock
  from typing import Optional, Tuple, List
  from ..utils.registry import Registry
  from ..attention_cell import MultiHeadAttentionCell, gen_self_attn_mask, gen_mem_attn_mask
  from ..layers import PositionalEmbedding, PositionwiseFFN, InitializerType
  from ..utils.config import CfgNode
  from ..sequence_sampler import BaseStepDecoder
label_desc: Use image node_id 3 for calling the TransformerDecoder obj's underlying member method code with example usage: obj.forward(data, valid_length, mem_data, mem_valid_length) and returns: out