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
  gated_proj
  pre_norm
      Whether to apply normalization before the attention layer
  use_qkv_bias
      Whether to use bias for both self attention and contextual attention
  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","gated_proj","pre_norm","Whether","to","apply","normalization","before","the","attention","layer","use_qkv_bias","Whether","to","use","bias","for","both","self","attention","and","contextual","attention","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",
      gated_proj: bool = False,
      pre_norm: bool = False,
      use_qkv_bias: bool = True,
      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
      gated_proj
      pre_norm
          Whether to apply normalization before the attention layer
      use_qkv_bias
          Whether to use bias for both self attention and contextual attention
      layout
          Layout of the input
      """
      super().__init__()
      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._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.Linear(
          out_features=3 * units, in_features=units, bias=use_qkv_bias
      )
      self.self_attention = MultiHeadAttentionCell(
          query_units=units,
          num_heads=num_heads,
          attention_dropout=self._attention_dropout,
          layout=attention_layout,
      )
      self.proj_in = nn.Linear(
          out_features=units, in_features=units, bias=True
      )
      self.attn_inter_q = nn.Linear(
          out_features=units, in_features=units, bias=use_qkv_bias
      )
      self.attn_inter_k = nn.Linear(
          out_features=units, in_features=mem_units, bias=use_qkv_bias
      )
      self.attn_inter_v = nn.Linear(
          out_features=units, in_features=mem_units, bias=use_qkv_bias
      )
      self.inter_attention = MultiHeadAttentionCell(
          query_units=units,
          num_heads=num_heads,
          attention_dropout=self._attention_dropout,
          layout=attention_layout,
      )
      self.proj_inter = nn.Linear(
          in_features=units, out_features=units, bias=True
      )
      self.ln_in = nn.LayerNorm(
          eps=layer_norm_eps, normalized_shape=units
      )
      self.ln_inter = nn.LayerNorm(
          eps=layer_norm_eps, normalized_shape=units
      )
      self.ffn = PositionwiseFFN(
          units=units,
          hidden_size=hidden_size,
          dropout=dropout,
          activation_dropout=activation_dropout,
          layer_norm_eps=layer_norm_eps,
          activation=activation,
          gated_proj=gated_proj,
          pre_norm=pre_norm,
      )
      self.init_weights()
  
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 1 to create a new TransformerDecoderLayer object from inherited base classes: nn with example: obj = TransformerDecoderLayer(units, mem_units, hidden_size, num_heads, activation_dropout, dropout, attention_dropout, layer_norm_eps, activation, gated_proj, pre_norm, use_qkv_bias, layout)