docstr: Parameters
  ----------
  units
  mem_units
      The number of units in the memory. By default, it is initialized to be the
      same as the units.
  hidden_size
  num_heads
  activation_dropout
  dropout
  attention_dropout
  layer_norm_eps
  activation
  pre_norm
      Whether to apply normalization before the attention layer
  use_qkv_bias
      Whether to use bias for both self attention and contextual attention
  weight_initializer
  bias_initializer
  dtype
      Data type
  layout
      Layout of the input
docstr_tok: ["Parameters","--","--","--","--","--","units","mem_units","The","number","of","units","in","the","memory",".","By","default",",","it","is","initialized","to","be","the","same","as","the","units",".","hidden_size","num_heads","activation_dropout","dropout","attention_dropout","layer_norm_eps","activation","pre_norm","Whether","to","apply","normalization","before","the","attention","layer","use_qkv_bias","Whether","to","use","bias","for","both","self","attention","and","contextual","attention","weight_initializer","bias_initializer","dtype","Data","type","layout","Layout","of","the","input"]
code: def __init__(
      self,
      units: int = 512,
      mem_units: Optional[int] = None,
      hidden_size: int = 2048,
      num_heads: int = 8,
      activation_dropout: float = 0.0,
      dropout: float = 0.1,
      attention_dropout: float = 0.1,
      layer_norm_eps: float = 1e-5,
      activation: str = "relu",
      pre_norm: bool = False,
      use_qkv_bias: bool = True,
      weight_initializer=None,
      bias_initializer="zeros",
      dtype="float32",
      layout="NT",
  ):
      """
  
      Parameters
      ----------
      units
      mem_units
          The number of units in the memory. By default, it is initialized to be the
          same as the units.
      hidden_size
      num_heads
      activation_dropout
      dropout
      attention_dropout
      layer_norm_eps
      activation
      pre_norm
          Whether to apply normalization before the attention layer
      use_qkv_bias
          Whether to use bias for both self attention and contextual attention
      weight_initializer
      bias_initializer
      dtype
          Data type
      layout
          Layout of the input
      """
      super().__init__()
      self._dtype = dtype
      self._units = units
      if mem_units is None:
          mem_units = units
      self._mem_units = mem_units
      self._pre_norm = pre_norm
      self._num_heads = num_heads
      self._attention_dropout = attention_dropout
      self._dtype = dtype
      self._layout = layout
      assert layout in ["TN", "NT"], (
          "Invalid layout received = {}. "
          'Only "TN" and "NT" are accepted!'.format(layout)
      )
      attention_layout = "NTK" if layout == "NT" else "TNK"
      self.dropout_layer = nn.Dropout(dropout)
      if units % num_heads:
          raise ValueError(
              "In Transformer, units should be divided exactly by the number of "
              "heads. Received units={}, num_heads={}".format(
                  units, num_heads
              )
          )
      self.attn_in_qkv = nn.Dense(
          3 * units,
          in_units=units,
          use_bias=use_qkv_bias,
          flatten=False,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          dtype=dtype,
      )
      self.self_attention = MultiHeadAttentionCell(
          query_units=units,
          num_heads=num_heads,
          attention_dropout=self._attention_dropout,
          dtype=dtype,
          layout=attention_layout,
      )
      self.proj_in = nn.Dense(
          units=units,
          in_units=units,
          flatten=False,
          use_bias=True,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          dtype=dtype,
      )
      self.attn_inter_q = nn.Dense(
          units,
          in_units=units,
          use_bias=use_qkv_bias,
          flatten=False,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          dtype=dtype,
      )
      self.attn_inter_k = nn.Dense(
          units,
          in_units=mem_units,
          use_bias=use_qkv_bias,
          flatten=False,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          dtype=dtype,
      )
      self.attn_inter_v = nn.Dense(
          units,
          in_units=mem_units,
          use_bias=use_qkv_bias,
          flatten=False,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          dtype=dtype,
      )
      self.inter_attention = MultiHeadAttentionCell(
          query_units=units,
          num_heads=num_heads,
          attention_dropout=self._attention_dropout,
          dtype=dtype,
          layout=attention_layout,
      )
      self.proj_inter = nn.Dense(
          units=units,
          in_units=units,
          flatten=False,
          use_bias=True,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          dtype=dtype,
      )
      # TODO(sxjscience) Add DType to LayerNorm
      self.ln_in = nn.LayerNorm(
          epsilon=layer_norm_eps, in_channels=units
      )
      self.ln_inter = nn.LayerNorm(
          epsilon=layer_norm_eps, in_channels=units
      )
      self.ffn = PositionwiseFFN(
          units=units,
          hidden_size=hidden_size,
          dropout=dropout,
          activation_dropout=activation_dropout,
          weight_initializer=weight_initializer,
          bias_initializer=bias_initializer,
          layer_norm_eps=layer_norm_eps,
          activation=activation,
          pre_norm=pre_norm,
          dtype=dtype,
      )
  
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 1 to create a new TransformerDecoderLayer object from inherited base classes: HybridBlock with example: obj = TransformerDecoderLayer(units, mem_units, hidden_size, num_heads, activation_dropout, dropout, attention_dropout, layer_norm_eps, activation, pre_norm, use_qkv_bias, weight_initializer, bias_initializer, dtype, layout)