class_docstr: An Apple Neural Engine (ANE) optimized Transformer implementation
      
docstr: Notation:
      src_seq_len:        The sequence length of the source sequence which is the input to the TransformerEncoder
      tgt_seq_len:        The sequence length of the target sequence which is the input to the TransformerDecoder
  
  Args:
      encoder_input:      Float tensor input to the TransformerEncoder
      decoder_input:      Float tensor input to the TransformerDecoder
      encoder_pos_embed:  Same shape and dtype as `encoder_input`. Serves as additive positional encodings to `encoder_input`
      decoder_pos_embed:  Same shape and dtype as `decoder_input`. Serves as additive positional encodings to `decoder_input`
      encoder_k_mask:     Float tensor similar to the `src_key_padding_mask` in `torch.nn.Transformer.forward`. Example use: masking zero-padded tokens in the source sequence
      decoder_k_mask:     Float tensor similar to the `tgt_key_padding_mask` in `torch.nn.Transformer.forward`. Example use: masking zero-padded tokens in the target sequence
      encoder_qk_mask:    Float tensor similar to `src_mask` in `torch.nn.Transformer.forward`. Example use: masking future tokens in the encoder self-attention
      decoder_qk_mask:    Float tensor similar to `tgt_mask` in `torch.nn.Transformer.forward`. Example use: masking future tokens in the decoder self-attention
  
  Shapes:
      encoder_input:      (batch_size, embed_dim, 1, src_seq_len)
      decoder_input:      (batch_size, embed_dim, 1, tgt_seq_len)
      encoder_pos_embed:  (batch_size, embed_dim, 1, src_seq_len)
      decoder_pos_embed:  (batch_size, embed_dim, 1, tgt_seq_len)
      encoder_k_mask:     (batch_size, src_seq_len, 1, 1)
      decoder_k_mask:     (batch_size, tgt_seq_len, 1, 1)
      encoder_qk_mask:    (batch_size, src_seq_len, 1, src_seq_len)
      decoder_qk_mask:    (batch_size, tgt_seq_len, 1, tgt_seq_len)
  
  Returns:
      decoder_output:     Output of the TransformerDecoder
      encoder_output:     Output of the TransformerEncoder
  
  
  Note: All arguments ending in "_mask", are applied additively on the intermediate tensor right before softmax in the attention function.
  The recommended float value for preventing attention is -1e4. This allows for composition of multiple masks while staying in the float16-friendly range.
  Use a value of 0 to keep attention unchanged.
code: def forward(
      self,
      encoder_input,
      decoder_input,
      encoder_pos_embed=None,
      decoder_pos_embed=None,
      encoder_k_mask=None,
      decoder_k_mask=None,
      encoder_qk_mask=None,
      decoder_qk_mask=None,
  ):
      """
      Notation:
          src_seq_len:        The sequence length of the source sequence which is the input to the TransformerEncoder
          tgt_seq_len:        The sequence length of the target sequence which is the input to the TransformerDecoder
  
      Args:
          encoder_input:      Float tensor input to the TransformerEncoder
          decoder_input:      Float tensor input to the TransformerDecoder
          encoder_pos_embed:  Same shape and dtype as `encoder_input`. Serves as additive positional encodings to `encoder_input`
          decoder_pos_embed:  Same shape and dtype as `decoder_input`. Serves as additive positional encodings to `decoder_input`
          encoder_k_mask:     Float tensor similar to the `src_key_padding_mask` in `torch.nn.Transformer.forward`. Example use: masking zero-padded tokens in the source sequence
          decoder_k_mask:     Float tensor similar to the `tgt_key_padding_mask` in `torch.nn.Transformer.forward`. Example use: masking zero-padded tokens in the target sequence
          encoder_qk_mask:    Float tensor similar to `src_mask` in `torch.nn.Transformer.forward`. Example use: masking future tokens in the encoder self-attention
          decoder_qk_mask:    Float tensor similar to `tgt_mask` in `torch.nn.Transformer.forward`. Example use: masking future tokens in the decoder self-attention
  
      Shapes:
          encoder_input:      (batch_size, embed_dim, 1, src_seq_len)
          decoder_input:      (batch_size, embed_dim, 1, tgt_seq_len)
          encoder_pos_embed:  (batch_size, embed_dim, 1, src_seq_len)
          decoder_pos_embed:  (batch_size, embed_dim, 1, tgt_seq_len)
          encoder_k_mask:     (batch_size, src_seq_len, 1, 1)
          decoder_k_mask:     (batch_size, tgt_seq_len, 1, 1)
          encoder_qk_mask:    (batch_size, src_seq_len, 1, src_seq_len)
          decoder_qk_mask:    (batch_size, tgt_seq_len, 1, tgt_seq_len)
  
      Returns:
          decoder_output:     Output of the TransformerDecoder
          encoder_output:     Output of the TransformerEncoder
  
  
      Note: All arguments ending in "_mask", are applied additively on the intermediate tensor right before softmax in the attention function.
      The recommended float value for preventing attention is -1e4. This allows for composition of multiple masks while staying in the float16-friendly range.
      Use a value of 0 to keep attention unchanged.
      """
      # Verify ranks
      assert_rank(encoder_input, "encoder_input", 4)
      assert_rank(encoder_input, "decoder_input", 4)
      assert_rank(encoder_pos_embed, "encoder_pos_embed", 4)
      assert_rank(decoder_pos_embed, "decoder_pos_embed", 4)
      assert_rank(encoder_k_mask, "encoder_k_mask", 4)
      assert_rank(decoder_k_mask, "decoder_k_mask", 4)
      assert_rank(encoder_k_mask, "encoder_qk_mask", 4)
      assert_rank(decoder_k_mask, "decoder_qk_mask", 4)
  
      # Verify and prepare encoder inputs
      batch_size, _, _, src_seq_len = encoder_input.shape
      assert_shape(
          encoder_input,
          "encoder_input",
          [batch_size, self.embed_dim, 1, src_seq_len],
      )
  
      if encoder_pos_embed is not None:
          assert_shape(
              encoder_pos_embed,
              "encoder_pos_embed",
              [batch_size, self.embed_dim, 1, src_seq_len],
          )
      if encoder_k_mask is not None:
          assert_shape(
              encoder_k_mask,
              "encoder_k_mask",
              [batch_size, src_seq_len, 1, 1],
          )
      if encoder_qk_mask is not None:
          assert_shape(
              encoder_qk_mask,
              "encoder_qk_mask",
              [batch_size, src_seq_len, 1, src_seq_len],
          )
  
      # Verify and prepare decoder inputs
      batch_size, _, _, tgt_seq_len = decoder_input.shape
      assert_shape(
          decoder_input,
          "decoder_input",
          [batch_size, self.embed_dim, 1, tgt_seq_len],
      )
  
      if decoder_pos_embed is not None:
          assert_shape(
              decoder_pos_embed,
              "decoder_pos_embed",
              [batch_size, self.embed_dim, 1, tgt_seq_len],
          )
      if decoder_k_mask is not None:
          assert_shape(
              decoder_k_mask,
              "decoder_k_mask",
              [batch_size, tgt_seq_len, 1, 1],
          )
      if decoder_qk_mask is not None:
          assert_shape(
              decoder_qk_mask,
              "decoder_qk_mask",
              [batch_size, tgt_seq_len, 1, tgt_seq_len],
          )
  
      # TransformerEncoder forward pass
      encoder_output = self.encoder(
          encoder_input,
          pos_embed=encoder_pos_embed,
          qk_mask=encoder_qk_mask,
          k_mask=encoder_k_mask,
      )
  
      # TransformerDecoder forward pass
      decoder_output = self.decoder(
          decoder_input,
          encoder_output,
          encoder_k_mask=encoder_k_mask,
          decoder_k_mask=decoder_k_mask,
          decoder_qk_mask=decoder_qk_mask,
          decoder_pos_embed=decoder_pos_embed,
          encoder_pos_embed=encoder_pos_embed,
          return_intermediate=self.return_intermediate_decoder,
      )
  
      return decoder_output, encoder_output
  
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 2 for calling the AppleNeuralEngineTransformer obj's underlying member method code with example usage: obj.forward(encoder_input, decoder_input, encoder_pos_embed, decoder_pos_embed, encoder_k_mask, decoder_k_mask, encoder_qk_mask, decoder_qk_mask) and returns: decoder_output,  encoder_output