class_docstr: An Apple Neural Engine (ANE) optimized Transformer implementation
      
docstr: Args:
      embed_dim:                          Dimensionality of the embedding space attention is computed in
      ffn_dim:                            Number of channels to use in the feed-forward network's hidden layer
      enc_self_attn_type:                 The self-attention module to use in each layer of Transformer encoder
      enc_ffn_type:                       The feed-forward network module to use in each layer of the Transformer encoder
      dec_self_attn_type:                 The self-attention module to use in each layer of Transformer decoder
      enc_dec_attn_type:                  The encoder-decoder attention module to use in each layer of Transformer decoder
      dec_ffn_type:                       The feed-forward network module to use in each layer of the Transformer decoder
      nb_enc_layers:                      Number of identically configured Transformer encoder layers to stack
      nb_dec_layers:                      Number of identically configured Transformer decoder layers to stack
      nb_attention_heads:                 Number of attention heads generate in each and every attention block in the Transformer model
      dropout:                            The dropout probability (`1- keep_probability`) to use on both the attention weights and the
                                          output of the attention block in each and every attention block in the Transformer model
      return_intermediate_decoder_layers: If True, returns the output of all Transformer decoder layers stacked in the 0-th axis.
                                          Example use case: When supervising the Transformer on all intermediate outputs for training stability
  
  Note: The positional embeddings are passed by the caller of forward() and are not part of this module
  
  Note: The default configuration reflects the "base" configuration in the original Transformer paper [1] (page 9, table 3, row 1)
  
  [1] https://arxiv.org/pdf/1706.03762
code: def __init__(
      self,
      *,
      embed_dim=512,
      ffn_dim=2048,
      enc_self_attn_type=multihead_attention.ResidualSelfAttention,
      enc_ffn_type=ffn.ResidualFFN,
      dec_self_attn_type=multihead_attention.ResidualSelfAttention,
      enc_dec_attn_type=multihead_attention.ResidualMultiHeadAttention,
      dec_ffn_type=ffn.ResidualFFN,
      nb_enc_layers=6,
      nb_dec_layers=6,
      nb_attention_heads=8,
      dropout=0.1,
      return_intermediate_decoder_layers=False,
      **kwargs,
  ):
      """
      Args:
          embed_dim:                          Dimensionality of the embedding space attention is computed in
          ffn_dim:                            Number of channels to use in the feed-forward network's hidden layer
          enc_self_attn_type:                 The self-attention module to use in each layer of Transformer encoder
          enc_ffn_type:                       The feed-forward network module to use in each layer of the Transformer encoder
          dec_self_attn_type:                 The self-attention module to use in each layer of Transformer decoder
          enc_dec_attn_type:                  The encoder-decoder attention module to use in each layer of Transformer decoder
          dec_ffn_type:                       The feed-forward network module to use in each layer of the Transformer decoder
          nb_enc_layers:                      Number of identically configured Transformer encoder layers to stack
          nb_dec_layers:                      Number of identically configured Transformer decoder layers to stack
          nb_attention_heads:                 Number of attention heads generate in each and every attention block in the Transformer model
          dropout:                            The dropout probability (`1- keep_probability`) to use on both the attention weights and the
                                              output of the attention block in each and every attention block in the Transformer model
          return_intermediate_decoder_layers: If True, returns the output of all Transformer decoder layers stacked in the 0-th axis.
                                              Example use case: When supervising the Transformer on all intermediate outputs for training stability
  
      Note: The positional embeddings are passed by the caller of forward() and are not part of this module
  
      Note: The default configuration reflects the "base" configuration in the original Transformer paper [1] (page 9, table 3, row 1)
  
      [1] https://arxiv.org/pdf/1706.03762
      """
      super().__init__()
      self.encoder = encoder.TransformerEncoder(
          layer=encoder.TransformerEncoderLayer(
              embed_dim,
              ffn_dim,
              enc_self_attn_type,
              enc_ffn_type,
              nb_attention_heads,
              dropout,
          ),
          num_layers=nb_enc_layers,
      )
      self.decoder = decoder.TransformerDecoder(
          layer=decoder.TransformerDecoderLayer(
              embed_dim,
              ffn_dim,
              dec_self_attn_type,
              enc_dec_attn_type,
              dec_ffn_type,
              nb_attention_heads,
              dropout,
          ),
          num_layers=nb_dec_layers,
      )
  
      self.return_intermediate_decoder = (
          return_intermediate_decoder_layers
      )
      self.embed_dim = embed_dim
  
file_path: ml-ane-transformers/ane_transformers/reference/transformer.py
imports: from ane_transformers.testing_utils import assert_rank, assert_shape
  import logging
  import torch
  import torch.nn
  from .None import encoder, decoder, multihead_attention, ffn
label_desc: Use image node_id 1 to create a new AppleNeuralEngineTransformer object from inherited base classes: nn with example: obj = AppleNeuralEngineTransformer()