class_docstr: Based on https://arxiv.org/abs/1805.00631
class_docstr_tok: ["Based","on","https",":","\/\/arxiv.org\/abs\/1805.00631"]
docstr: Args:
      x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
      encoder_padding_mask (ByteTensor, optional): binary
          ByteTensor of shape `(batch, src_len)` where padding
          elements are indicated by ``1``.
      need_attn (bool, optional): return attention weights
      need_head_weights (bool, optional): return attention weights
          for each head (default: return average over heads).
  
  Returns:
      encoded output of shape `(seq_len, batch, embed_dim)`
  
  The following are used for export tracing:
      prev_self_attn_state: [prev_sum, prev_pos]
          assumes AverageAttention without mask trick
      prev_attn_state: [prev_key, prev_value]
docstr_tok: ["Args",":","x","(","Tensor",")",":","input","to","the","layer","of","shape","`","(","seq_len",",","batch",",","embed_dim",")","`","encoder_padding_mask","(","ByteTensor",",","optional",")",":","binary","ByteTensor","of","shape","`","(","batch",",","src_len",")","`","where","padding","elements","are","indicated","by","``","1","``",".","need_attn","(","bool",",","optional",")",":","return","attention","weights","need_head_weights","(","bool",",","optional",")",":","return","attention","weights","for","each","head","(","default",":","return","average","over","heads",")",".","Returns",":","encoded","output","of","shape","`","(","seq_len",",","batch",",","embed_dim",")","`","The","following","are","used","for","export","tracing",":","prev_self_attn_state",":","[","prev_sum",",","prev_pos","]","assumes","AverageAttention","without","mask","trick","prev_attn_state",":","[","prev_key",",","prev_value","]"]
code: def forward(
      self,
      x,
      encoder_out=None,
      encoder_padding_mask=None,
      incremental_state=None,
      prev_self_attn_state=None,
      prev_attn_state=None,
      self_attn_mask=None,
      self_attn_padding_mask=None,
      need_attn=False,
      need_head_weights=False,
  ):
      """
      Args:
          x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
          encoder_padding_mask (ByteTensor, optional): binary
              ByteTensor of shape `(batch, src_len)` where padding
              elements are indicated by ``1``.
          need_attn (bool, optional): return attention weights
          need_head_weights (bool, optional): return attention weights
              for each head (default: return average over heads).
  
      Returns:
          encoded output of shape `(seq_len, batch, embed_dim)`
  
      The following are used for export tracing:
          prev_self_attn_state: [prev_sum, prev_pos]
              assumes AverageAttention without mask trick
          prev_attn_state: [prev_key, prev_value]
      """
      if need_head_weights:
          need_attn = True
  
      residual = x
      x = self.maybe_layer_norm(
          self.avg_attn_layer_norm, x, before=True
      )
  
      if prev_self_attn_state is not None:
          if incremental_state is None:
              incremental_state = {}
          prev_sum, prev_pos = prev_self_attn_state
          # (batch, embed) -> (seq, batch, embed)
          prev_sum = prev_sum.unsqueeze(0)
          saved_state = {"prev_sum": prev_sum, "prev_pos": prev_pos}
          self.avg_attn._set_input_buffer(
              incremental_state, saved_state
          )
  
      x, _ = self.avg_attn(
          value=x,
          mask_future_timesteps=True,
          incremental_state=incremental_state,
          mask_trick=self.training,
      )
  
      # differently than original paper, we use a single gate
      gate = torch.sigmoid(
          self.aan_gating_fc(torch.cat([residual, x], dim=-1))
      )
      x = gate * x + (1 - gate) * residual
  
      x = F.dropout(x, p=self.dropout, training=self.training)
      x = residual + x
      x = self.maybe_layer_norm(self.avg_attn_layer_norm, x, after=True)
  
      if self.encoder_attn is not None:
          residual = x
          x = self.maybe_layer_norm(
              self.encoder_attn_layer_norm, x, before=True
          )
          if prev_attn_state is not None:
              if incremental_state is None:
                  incremental_state = {}
              prev_key, prev_value = prev_attn_state[:2]
              saved_state = {
                  "prev_key": prev_key,
                  "prev_value": prev_value,
              }
              if len(prev_attn_state) >= 3:
                  saved_state[
                      "prev_key_padding_mask"
                  ] = prev_attn_state[2]
              self.encoder_attn._set_input_buffer(
                  incremental_state, saved_state
              )
  
          x, attn = self.encoder_attn(
              query=x,
              key=encoder_out,
              value=encoder_out,
              key_padding_mask=encoder_padding_mask,
              incremental_state=incremental_state,
              static_kv=True,
              need_weights=need_attn
              or (not self.training and self.need_attn),
              need_head_weights=need_head_weights,
          )
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          x = self.maybe_layer_norm(
              self.encoder_attn_layer_norm, x, after=True
          )
  
      residual = x
      x = self.maybe_layer_norm(self.final_layer_norm, x, before=True)
      x = self.activation_fn(self.fc1(x))
      x = F.dropout(
          x, p=self.activation_dropout, training=self.training
      )
      x = self.fc2(x)
      x = F.dropout(x, p=self.dropout, training=self.training)
      x = residual + x
      x = self.maybe_layer_norm(self.final_layer_norm, x, after=True)
  
      if self.onnx_trace and incremental_state is not None:
          saved_state = self.avg_attn._get_input_buffer(
              incremental_state
          )
          # remove sequence axis for export
          prev_sum = saved_state["prev_sum"]
          # (seq, batch, embed) -> (batch, embed)
          prev_sum = prev_sum.squeeze(0)
          prev_pos = saved_state["prev_pos"]
          self_attn_state = prev_sum, prev_pos
          return x, attn, self_attn_state
  
      return x, attn, None
  
imports: import math
  import random
  import torch
  import torch.nn
  import torch.nn.functional
  from fairseq import options, utils
  from fairseq.models import FairseqEncoder, FairseqEncoderDecoderModel, FairseqIncrementalDecoder, register_model, register_model_architecture, transformer
  from fairseq.modules import AdaptiveSoftmax, LayerNorm, MultiheadAttention, SinusoidalPositionalEmbedding
  from pytorch_translate import utils, vocab_reduction
  from pytorch_translate.average_attention import AverageAttention
  from pytorch_translate.common_layers import TransformerEmbedding, TransformerEncoderGivenEmbeddings, TransformerTokenEmbedding, VariableTracker
  from pytorch_translate.semi_supervised import SemiSupervisedModel
  from pytorch_translate.utils import torch_find
label_desc: Use image node_id 3 for calling the AANDecoderLayer obj's underlying member method code with example usage: obj.forward(x, encoder_out, encoder_padding_mask, incremental_state, prev__attn_state, prev_attn_state, _attn_mask, _attn_padding_mask, need_attn, need_head_weights) and returns: x,  attn,  None, x,  attn,  self_attn_state