docstr: Args:
      d_model (int): the dimension of keys/values/queries in
          :class:`MultiHeadedAttention`, also the input size of
          the first-layer of the :class:`PositionwiseFeedForward`.
      heads (int): the number of heads for MultiHeadedAttention.
      d_ff (int): the second-layer of the
          :class:`PositionwiseFeedForward`.
      dropout (float): dropout in residual, self-attn(dot) and
          feed-forward
      attention_dropout (float): dropout in context_attn  (and
          self-attn(avg))
      self_attn_type (string): type of self-attention scaled-dot,
          flash-scaled-dot, average
      max_relative_positions (int):
          Max distance between inputs in relative positions
          representations
      relative_positions_buckets (int):
          relative position bias see
          https://github.com/google-research/text-to-text-transfer-transformer
      aan_useffn (bool): Turn on the FFN layer in the AAN decoder
      full_context_alignment (bool):
          whether enable an extra full context decoder forward for
          alignment
      alignment_heads (int):
          N. of cross attention heads to use for alignment guiding
      pos_ffn_activation_fn (ActivationFunction):
          activation function choice for PositionwiseFeedForward layer
      add_qkvbias (bool): whether to add bias to the Key/Value nn.Linear
      num_kv (int): number of heads for KV when different vs Q (multiquery)
      add_ffnbias (bool): whether to add bias to the FF nn.Linear
      parallel_residual (bool): Use parallel residual connections in each layer block, as used
          by the GPT-J and GPT-NeoX models
      shared_layer_norm (bool): When using parallel residual, share the input and post
          attention layer norms.
      layer_norm (string): type of layer normalization standard/rms
      norm_eps (float): layer norm epsilon
      use_ckpting (List): layers for which we checkpoint for backward
      parallel_gpu (int): Number of gpu for tensor parallelism
      sliding_window (int): Width of the band mask and KV cache (cf Mistral Model)
      rotary_interleave (bool): Interleave the head dimensions when rotary
          embeddings are applied
      rotary_theta (int): rotary base theta
      rotary_dim (int): in some cases the rotary dim is lower than head dim
      num_experts (int): Number of experts for MoE
      num_experts_per_tok (int): Number of experts choice per token
docstr_tok: ["Args",":","d_model","(","int",")",":","the","dimension","of","keys\/values\/queries","in",":","class",":","`","MultiHeadedAttention","`",",","also","the","input","size","of","the","first-layer","of","the",":","class",":","`","PositionwiseFeedForward","`",".","heads","(","int",")",":","the","number","of","heads","for","MultiHeadedAttention",".","d_ff","(","int",")",":","the","second-layer","of","the",":","class",":","`","PositionwiseFeedForward","`",".","dropout","(","float",")",":","dropout","in","residual",",","self-attn","(","dot",")","and","feed-forward","attention_dropout","(","float",")",":","dropout","in","context_attn","(","and","self-attn","(","avg",")",")","self_attn_type","(","string",")",":","type","of","self-attention","scaled-dot",",","flash-scaled-dot",",","average","max_relative_positions","(","int",")",":","Max","distance","between","inputs","in","relative","positions","representations","relative_positions_buckets","(","int",")",":","relative","position","bias","see","https",":","\/\/github.com\/google-research\/text-to-text-transfer-transformer","aan_useffn","(","bool",")",":","Turn","on","the","FFN","layer","in","the","AAN","decoder","full_context_alignment","(","bool",")",":","whether","enable","an","extra","full","context","decoder","forward","for","alignment","alignment_heads","(","int",")",":","N.","of","cross","attention","heads","to","use","for","alignment","guiding","pos_ffn_activation_fn","(","ActivationFunction",")",":","activation","function","choice","for","PositionwiseFeedForward","layer","add_qkvbias","(","bool",")",":","whether","to","add","bias","to","the","Key\/Value","nn.Linear","num_kv","(","int",")",":","number","of","heads","for","KV","when","different","vs","Q","(","multiquery",")","add_ffnbias","(","bool",")",":","whether","to","add","bias","to","the","FF","nn.Linear","parallel_residual","(","bool",")",":","Use","parallel","residual","connections","in","each","layer","block",",","as","used","by","the","GPT-J","and","GPT-NeoX","models","shared_layer_norm","(","bool",")",":","When","using","parallel","residual",",","share","the","input","and","post","attention","layer","norms",".","layer_norm","(","string",")",":","type","of","layer","normalization","standard\/rms","norm_eps","(","float",")",":","layer","norm","epsilon","use_ckpting","(","List",")",":","layers","for","which","we","checkpoint","for","backward","parallel_gpu","(","int",")",":","Number","of","gpu","for","tensor","parallelism","sliding_window","(","int",")",":","Width","of","the","band","mask","and","KV","cache","(","cf","Mistral","Model",")","rotary_interleave","(","bool",")",":","Interleave","the","head","dimensions","when","rotary","embeddings","are","applied","rotary_theta","(","int",")",":","rotary","base","theta","rotary_dim","(","int",")",":","in","some","cases","the","rotary","dim","is","lower","than","head","dim","num_experts","(","int",")",":","Number","of","experts","for","MoE","num_experts_per_tok","(","int",")",":","Number","of","experts","choice","per","token"]
code: def __init__(
      self,
      d_model,
      heads,
      d_ff,
      dropout,
      attention_dropout,
      self_attn_type="scaled_dot",
      max_relative_positions=0,
      relative_positions_buckets=0,
      aan_useffn=False,
      full_context_alignment=False,
      alignment_heads=0,
      pos_ffn_activation_fn=ActivationFunction.relu,
      add_qkvbias=False,
      num_kv=0,
      add_ffnbias=True,
      parallel_residual=False,
      shared_layer_norm=False,
      layer_norm="standard",
      norm_eps=1e-6,
      use_ckpting=[],
      parallel_gpu=1,
      sliding_window=0,
      rotary_interleave=True,
      rotary_theta=1e4,
      rotary_dim=0,
      num_experts=0,
      num_experts_per_tok=2,
  ):
      """
      Args:
          d_model (int): the dimension of keys/values/queries in
              :class:`MultiHeadedAttention`, also the input size of
              the first-layer of the :class:`PositionwiseFeedForward`.
          heads (int): the number of heads for MultiHeadedAttention.
          d_ff (int): the second-layer of the
              :class:`PositionwiseFeedForward`.
          dropout (float): dropout in residual, self-attn(dot) and
              feed-forward
          attention_dropout (float): dropout in context_attn  (and
              self-attn(avg))
          self_attn_type (string): type of self-attention scaled-dot,
              flash-scaled-dot, average
          max_relative_positions (int):
              Max distance between inputs in relative positions
              representations
          relative_positions_buckets (int):
              relative position bias see
              https://github.com/google-research/text-to-text-transfer-transformer
          aan_useffn (bool): Turn on the FFN layer in the AAN decoder
          full_context_alignment (bool):
              whether enable an extra full context decoder forward for
              alignment
          alignment_heads (int):
              N. of cross attention heads to use for alignment guiding
          pos_ffn_activation_fn (ActivationFunction):
              activation function choice for PositionwiseFeedForward layer
          add_qkvbias (bool): whether to add bias to the Key/Value nn.Linear
          num_kv (int): number of heads for KV when different vs Q (multiquery)
          add_ffnbias (bool): whether to add bias to the FF nn.Linear
          parallel_residual (bool): Use parallel residual connections in each layer block, as used
              by the GPT-J and GPT-NeoX models
          shared_layer_norm (bool): When using parallel residual, share the input and post
              attention layer norms.
          layer_norm (string): type of layer normalization standard/rms
          norm_eps (float): layer norm epsilon
          use_ckpting (List): layers for which we checkpoint for backward
          parallel_gpu (int): Number of gpu for tensor parallelism
          sliding_window (int): Width of the band mask and KV cache (cf Mistral Model)
          rotary_interleave (bool): Interleave the head dimensions when rotary
              embeddings are applied
          rotary_theta (int): rotary base theta
          rotary_dim (int): in some cases the rotary dim is lower than head dim
          num_experts (int): Number of experts for MoE
          num_experts_per_tok (int): Number of experts choice per token
      """
      super(TransformerDecoderLayerBase, self).__init__()
  
      if self_attn_type in ["scaled-dot", "scaled-dot-flash"]:
          self.self_attn = MultiHeadedAttention(
              heads,
              d_model,
              dropout=attention_dropout,
              max_relative_positions=max_relative_positions,
              relative_positions_buckets=relative_positions_buckets,
              rotary_interleave=rotary_interleave,
              rotary_theta=rotary_theta,
              rotary_dim=rotary_dim,
              attn_type="self",
              self_attn_type=self_attn_type,
              add_qkvbias=add_qkvbias,
              num_kv=num_kv,
              use_ckpting=use_ckpting,
              parallel_gpu=parallel_gpu,
          )
      elif self_attn_type == "average":
          self.self_attn = AverageAttention(
              d_model, dropout=attention_dropout, aan_useffn=aan_useffn
          )
  
      if num_experts > 0:
          self.feed_forward = MoE(
              num_experts,
              num_experts_per_tok,
              d_model,
              d_ff,
              dropout,
              pos_ffn_activation_fn,
              add_ffnbias,
              parallel_residual,
              layer_norm,
              norm_eps,
              use_ckpting=use_ckpting,
              parallel_gpu=parallel_gpu,
          )
      else:
          self.feed_forward = PositionwiseFeedForward(
              d_model,
              d_ff,
              dropout,
              pos_ffn_activation_fn,
              add_ffnbias,
              parallel_residual,
              layer_norm,
              norm_eps,
              use_ckpting=use_ckpting,
              parallel_gpu=parallel_gpu,
          )
      self.parallel_residual = parallel_residual
      self.shared_layer_norm = shared_layer_norm
      if layer_norm == "standard":
          self.layer_norm_1 = nn.LayerNorm(d_model, eps=norm_eps)
          if parallel_residual and not shared_layer_norm:
              self.layer_norm_res = nn.LayerNorm(d_model, eps=norm_eps)
      elif layer_norm == "rms":
          self.layer_norm_1 = RMSNorm(d_model, eps=norm_eps)
          if parallel_residual and not shared_layer_norm:
              self.layer_norm_res = RMSNorm(d_model, eps=norm_eps)
      else:
          raise ValueError(
              f"{layer_norm} layer norm type is not supported"
          )
  
      self.dropout = nn.Dropout(dropout)
      self.dropout_p = dropout
      self.full_context_alignment = full_context_alignment
      self.alignment_heads = alignment_heads
      self.sliding_window = sliding_window
      self.self_attn_type = self_attn_type
  
imports: import torch
  import torch.nn
  from onmt.decoders.decoder import DecoderBase
  from onmt.modules import MultiHeadedAttention, AverageAttention
  from onmt.modules.position_ffn import PositionwiseFeedForward
  from onmt.modules.position_ffn import ActivationFunction
  from onmt.modules.moe import MoE
  from onmt.utils.misc import sequence_mask
  from onmt.modules.rmsnorm import RMSNorm
label_desc: Use image node_id 1 to create a new TransformerDecoderLayerBase object from inherited base classes: nn with example: obj = TransformerDecoderLayerBase(d_model, heads, d_ff, dropout, attention_dropout, _attn_type, max_relative_positions, relative_positions_buckets, aan_useffn, full_context_alignment, alignment_heads, pos_ffn_activation_fn, add_qkvbias, num_kv, add_ffnbias, parallel_residual, shared_layer_norm, layer_norm, norm_eps, use_ckpting, parallel_gpu, sliding_window, rotary_interleave, rotary_theta, rotary_dim, num_experts, num_experts_per_tok)