class_docstr: Average Attention module from
  "Accelerating Neural Transformer via an Average Attention Network"
  :cite:`DBLP:journals/corr/abs-1805-00631`.
  
  Args:
     model_dim (int): the dimension of keys/values/queries,
         must be divisible by head_count
     dropout (float): dropout parameter
     pos_ffn_activation_fn (ActivationFunction):
         activation function choice for PositionwiseFeedForward layer
docstr: Args:
      layer_in (FloatTensor): ``(batch, t_len, dim)``
  
  Returns:
      (FloatTensor, FloatTensor):
  
      * gating_out ``(batch, tlen, dim)``
      * average_out average attention
          ``(batch, input_len, dim)``
code: def forward(self, layer_in, mask=None, step=None):
      """
      Args:
          layer_in (FloatTensor): ``(batch, t_len, dim)``
  
      Returns:
          (FloatTensor, FloatTensor):
  
          * gating_out ``(batch, tlen, dim)``
          * average_out average attention
              ``(batch, input_len, dim)``
      """
  
      batch_size = layer_in.size(0)
      t_len = layer_in.size(1)
      mask = (
          cumulative_average_mask(batch_size, t_len, layer_in.device)
          if not self.layer_cache[0]
          else None
      )
      average_out = cumulative_average(
          layer_in, self.layer_cache, mask, step
      )
      if self.aan_useffn:
          average_out = self.average_layer(average_out)
      gating_out = self.gating_layer(
          torch.cat((layer_in, average_out), -1)
      )
      input_gate, forget_gate = torch.chunk(gating_out, 2, dim=2)
      gating_out = (
          torch.sigmoid(input_gate) * layer_in
          + torch.sigmoid(forget_gate) * average_out
      )
  
      return gating_out, average_out
  
file_path: OpenNMT-py/onmt/modules/average_attn.py
imports: import torch
  import torch.nn
  from torch import Tensor
  from typing import Optional
  from onmt.modules.position_ffn import PositionwiseFeedForward
  from onmt.modules.position_ffn import ActivationFunction
label_desc: Use image node_id 2 for calling the AverageAttention obj's underlying member method code with example usage: obj.forward(layer_in, mask, step) and returns: gating_out,  average_out