prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  The comment in onmt.modules.average_attn may be wrong
  The comment in onmt.modules.average_attn.AverageAttention.forward says that it gets a input of size `[batch_size x input_len x model_dim]`, and outputs a gating_outputs and a average_outputs both with size `[batch_size x 1 x model_dim]`. However, in my test the two outputs are both with size [batch_size x input_len x model_dim].
  ```python
  import torch
  import onmt.modules.average_attn as Att
  model = Att.AverageAttention(3)
  a = torch.arange(24).reshape(2,4,3).float()
  print("intput size:", a.size())
  o1, o2 = model(a)
  print("o1 size:", o1.size())
  print("o2.size:", o2.size())
  ```
  intput size: torch.Size([2, 4, 3])
  o1 size: torch.Size([2, 4, 3])
  o2 size: torch.Size([2, 4, 3])
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `onmt/modules/average_attn.py`
  Content:
  ```
  1 # -*- coding: utf-8 -*-
  2 """ Average Attention module """
  3 
  4 import torch
  5 import torch.nn as nn
  6 
  7 from onmt.modules.position_ffn import PositionwiseFeedForward
  8 
  9 
  10 class AverageAttention(nn.Module):
  11     """
  12     Average Attention module from
  13     "Accelerating Neural Transformer via an Average Attention Network"
  14     :cite:`https://arxiv.org/abs/1805.00631`.
  15 
  16     Args:
  17        model_dim (int): the dimension of keys/values/queries,
  18            must be divisible by head_count
  19        dropout (float): dropout parameter
  20     """
  21 
  22     def __init__(self, model_dim, dropout=0.1):
  23         self.model_dim = model_dim
  24 
  25         super(AverageAttention, self).__init__()
  26 
  27         self.average_layer = PositionwiseFeedForward(model_dim, model_dim,
  28                                                      dropout)
  29         self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)
  30 
  31     def cumulative_average_mask(self, batch_size, inputs_len):
  32         """
  33         Builds the mask to compute the cumulative average as described in
  34         https://arxiv.org/abs/1805.00631 -- Figure 3
  35 
  36         Args:
  37             batch_size (int): batch size
  38             inputs_len (int): length of the inputs
  39 
  40         Returns:
  41             (`FloatTensor`):
  42 
  43             * A Tensor of shape `[batch_size x input_len x input_len]`
  44         """
  45 
  46         triangle = torch.tril(torch.ones(inputs_len, inputs_len))
  47         weights = torch.ones(1, inputs_len) / torch.arange(
  48             1, inputs_len + 1, dtype=torch.float)
  49         mask = triangle * weights.transpose(0, 1)
  50 
  51         return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)
  52 
  53     def cumulative_average(self, inputs, mask_or_step,
  54                            layer_cache=None, step=None):
  55         """
  56         Computes the cumulative average as described in
  57         https://arxiv.org/abs/1805.00631 -- Equations (1) (5) (6)
  58 
  59         Args:
  60             inputs (`FloatTensor`): sequence to average
  61                 `[batch_size x input_len x dimension]`
  62             mask_or_step: if cache is set, this is assumed
  63                 to be the current step of the
  64                 dynamic decoding. Otherwise, it is the mask matrix
  65                 used to compute the cumulative average.
  66             cache: a dictionary containing the cumulative average
  67                 of the previous step.
  68         """
  69         if layer_cache is not None:
  70             step = mask_or_step
  71             device = inputs.device
  72             average_attention = (inputs + step *
  73                                  layer_cache["prev_g"].to(device)) / (step + 1)
  74             layer_cache["prev_g"] = average_attention
  75             return average_attention
  76         else:
  77             mask = mask_or_step
  78             return torch.matmul(mask, inputs)
  79 
  80     def forward(self, inputs, mask=None, layer_cache=None, step=None):
  81         """
  82         Args:
  83             inputs (`FloatTensor`): `[batch_size x input_len x model_dim]`
  84 
  85         Returns:
  86             (`FloatTensor`, `FloatTensor`):
  87 
  88             * gating_outputs `[batch_size x 1 x model_dim]`
  89             * average_outputs average attention `[batch_size x 1 x model_dim]`
  90         """
  91         batch_size = inputs.size(0)
  92         inputs_len = inputs.size(1)
  93 
  94         device = inputs.device
  95         average_outputs = self.cumulative_average(
  96           inputs, self.cumulative_average_mask(batch_size,
  97                                                inputs_len).to(device).float()
  98           if layer_cache is None else step, layer_cache=layer_cache)
  99         average_outputs = self.average_layer(average_outputs)
  100         gating_outputs = self.gating_layer(torch.cat((inputs,
  101                                                       average_outputs), -1))
  102         input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)
  103         gating_outputs = torch.sigmoid(input_gate) * inputs + \
  104             torch.sigmoid(forget_gate) * average_outputs
  105 
  106         return gating_outputs, average_outputs
  107
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/onmt/modules/average_attn.py b/onmt/modules/average_attn.py
  --- a/onmt/modules/average_attn.py
  +++ b/onmt/modules/average_attn.py
  @@ -65,6 +65,9 @@
                   used to compute the cumulative average.
               cache: a dictionary containing the cumulative average
                   of the previous step.
  +
  +        Returns:
  +            a tensor of the same shape and type as :obj:`inputs`.
           """
           if layer_cache is not None:
               step = mask_or_step
  @@ -85,8 +88,9 @@
           Returns:
               (`FloatTensor`, `FloatTensor`):
   
  -            * gating_outputs `[batch_size x 1 x model_dim]`
  -            * average_outputs average attention `[batch_size x 1 x model_dim]`
  +            * gating_outputs `[batch_size x input_len x model_dim]`
  +            * average_outputs average attention
  +                `[batch_size x input_len x model_dim]`
           """
           batch_size = inputs.size(0)
           inputs_len = inputs.size(1)
  
verification_info: {"golden_diff": "diff --git a/onmt/modules/average_attn.py b/onmt/modules/average_attn.py\n--- a/onmt/modules/average_attn.py\n+++ b/onmt/modules/average_attn.py\n@@ -65,6 +65,9 @@\n                 used to compute the cumulative average.\n             cache: a dictionary containing the cumulative average\n                 of the previous step.\n+\n+        Returns:\n+            a tensor of the same shape and type as :obj:`inputs`.\n         \"\"\"\n         if layer_cache is not None:\n             step = mask_or_step\n@@ -85,8 +88,9 @@\n         Returns:\n             (`FloatTensor`, `FloatTensor`):\n \n-            * gating_outputs `[batch_size x 1 x model_dim]`\n-            * average_outputs average attention `[batch_size x 1 x model_dim]`\n+            * gating_outputs `[batch_size x input_len x model_dim]`\n+            * average_outputs average attention\n+                `[batch_size x input_len x model_dim]`\n         \"\"\"\n         batch_size = inputs.size(0)\n         inputs_len = inputs.size(1)\n", "issue": "The comment in onmt.modules.average_attn may be wrong\nThe comment in onmt.modules.average_attn.AverageAttention.forward says that it gets a input of size `[batch_size x input_len x model_dim]`, and outputs a gating_outputs and a average_outputs both with size `[batch_size x 1 x model_dim]`. However, in my test the two outputs are both with size [batch_size x input_len x model_dim].\r\n```python\r\nimport torch\r\nimport onmt.modules.average_attn as Att\r\nmodel = Att.AverageAttention(3)\r\na = torch.arange(24).reshape(2,4,3).float()\r\nprint(\"intput size:\", a.size())\r\no1, o2 = model(a)\r\nprint(\"o1 size:\", o1.size())\r\nprint(\"o2.size:\", o2.size())\r\n```\r\nintput size: torch.Size([2, 4, 3])\r\no1 size: torch.Size([2, 4, 3])\r\no2 size: torch.Size([2, 4, 3])\n", "before_files": [{"content": "# -*- coding: utf-8 -*-\n\"\"\" Average Attention module \"\"\"\n\nimport torch\nimport torch.nn as nn\n\nfrom onmt.modules.position_ffn import PositionwiseFeedForward\n\n\nclass AverageAttention(nn.Module):\n    \"\"\"\n    Average Attention module from\n    \"Accelerating Neural Transformer via an Average Attention Network\"\n    :cite:`https://arxiv.org/abs/1805.00631`.\n\n    Args:\n       model_dim (int): the dimension of keys/values/queries,\n           must be divisible by head_count\n       dropout (float): dropout parameter\n    \"\"\"\n\n    def __init__(self, model_dim, dropout=0.1):\n        self.model_dim = model_dim\n\n        super(AverageAttention, self).__init__()\n\n        self.average_layer = PositionwiseFeedForward(model_dim, model_dim,\n                                                     dropout)\n        self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)\n\n    def cumulative_average_mask(self, batch_size, inputs_len):\n        \"\"\"\n        Builds the mask to compute the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Figure 3\n\n        Args:\n            batch_size (int): batch size\n            inputs_len (int): length of the inputs\n\n        Returns:\n            (`FloatTensor`):\n\n            * A Tensor of shape `[batch_size x input_len x input_len]`\n        \"\"\"\n\n        triangle = torch.tril(torch.ones(inputs_len, inputs_len))\n        weights = torch.ones(1, inputs_len) / torch.arange(\n            1, inputs_len + 1, dtype=torch.float)\n        mask = triangle * weights.transpose(0, 1)\n\n        return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)\n\n    def cumulative_average(self, inputs, mask_or_step,\n                           layer_cache=None, step=None):\n        \"\"\"\n        Computes the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Equations (1) (5) (6)\n\n        Args:\n            inputs (`FloatTensor`): sequence to average\n                `[batch_size x input_len x dimension]`\n            mask_or_step: if cache is set, this is assumed\n                to be the current step of the\n                dynamic decoding. Otherwise, it is the mask matrix\n                used to compute the cumulative average.\n            cache: a dictionary containing the cumulative average\n                of the previous step.\n        \"\"\"\n        if layer_cache is not None:\n            step = mask_or_step\n            device = inputs.device\n            average_attention = (inputs + step *\n                                 layer_cache[\"prev_g\"].to(device)) / (step + 1)\n            layer_cache[\"prev_g\"] = average_attention\n            return average_attention\n        else:\n            mask = mask_or_step\n            return torch.matmul(mask, inputs)\n\n    def forward(self, inputs, mask=None, layer_cache=None, step=None):\n        \"\"\"\n        Args:\n            inputs (`FloatTensor`): `[batch_size x input_len x model_dim]`\n\n        Returns:\n            (`FloatTensor`, `FloatTensor`):\n\n            * gating_outputs `[batch_size x 1 x model_dim]`\n            * average_outputs average attention `[batch_size x 1 x model_dim]`\n        \"\"\"\n        batch_size = inputs.size(0)\n        inputs_len = inputs.size(1)\n\n        device = inputs.device\n        average_outputs = self.cumulative_average(\n          inputs, self.cumulative_average_mask(batch_size,\n                                               inputs_len).to(device).float()\n          if layer_cache is None else step, layer_cache=layer_cache)\n        average_outputs = self.average_layer(average_outputs)\n        gating_outputs = self.gating_layer(torch.cat((inputs,\n                                                      average_outputs), -1))\n        input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)\n        gating_outputs = torch.sigmoid(input_gate) * inputs + \\\n            torch.sigmoid(forget_gate) * average_outputs\n\n        return gating_outputs, average_outputs\n", "path": "onmt/modules/average_attn.py"}], "after_files": [{"content": "# -*- coding: utf-8 -*-\n\"\"\" Average Attention module \"\"\"\n\nimport torch\nimport torch.nn as nn\n\nfrom onmt.modules.position_ffn import PositionwiseFeedForward\n\n\nclass AverageAttention(nn.Module):\n    \"\"\"\n    Average Attention module from\n    \"Accelerating Neural Transformer via an Average Attention Network\"\n    :cite:`https://arxiv.org/abs/1805.00631`.\n\n    Args:\n       model_dim (int): the dimension of keys/values/queries,\n           must be divisible by head_count\n       dropout (float): dropout parameter\n    \"\"\"\n\n    def __init__(self, model_dim, dropout=0.1):\n        self.model_dim = model_dim\n\n        super(AverageAttention, self).__init__()\n\n        self.average_layer = PositionwiseFeedForward(model_dim, model_dim,\n                                                     dropout)\n        self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)\n\n    def cumulative_average_mask(self, batch_size, inputs_len):\n        \"\"\"\n        Builds the mask to compute the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Figure 3\n\n        Args:\n            batch_size (int): batch size\n            inputs_len (int): length of the inputs\n\n        Returns:\n            (`FloatTensor`):\n\n            * A Tensor of shape `[batch_size x input_len x input_len]`\n        \"\"\"\n\n        triangle = torch.tril(torch.ones(inputs_len, inputs_len))\n        weights = torch.ones(1, inputs_len) / torch.arange(\n            1, inputs_len + 1, dtype=torch.float)\n        mask = triangle * weights.transpose(0, 1)\n\n        return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)\n\n    def cumulative_average(self, inputs, mask_or_step,\n                           layer_cache=None, step=None):\n        \"\"\"\n        Computes the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Equations (1) (5) (6)\n\n        Args:\n            inputs (`FloatTensor`): sequence to average\n                `[batch_size x input_len x dimension]`\n            mask_or_step: if cache is set, this is assumed\n                to be the current step of the\n                dynamic decoding. Otherwise, it is the mask matrix\n                used to compute the cumulative average.\n            cache: a dictionary containing the cumulative average\n                of the previous step.\n\n        Returns:\n            a tensor of the same shape and type as :obj:`inputs`.\n        \"\"\"\n        if layer_cache is not None:\n            step = mask_or_step\n            device = inputs.device\n            average_attention = (inputs + step *\n                                 layer_cache[\"prev_g\"].to(device)) / (step + 1)\n            layer_cache[\"prev_g\"] = average_attention\n            return average_attention\n        else:\n            mask = mask_or_step\n            return torch.matmul(mask, inputs)\n\n    def forward(self, inputs, mask=None, layer_cache=None, step=None):\n        \"\"\"\n        Args:\n            inputs (`FloatTensor`): `[batch_size x input_len x model_dim]`\n\n        Returns:\n            (`FloatTensor`, `FloatTensor`):\n\n            * gating_outputs `[batch_size x input_len x model_dim]`\n            * average_outputs average attention\n                `[batch_size x input_len x model_dim]`\n        \"\"\"\n        batch_size = inputs.size(0)\n        inputs_len = inputs.size(1)\n\n        device = inputs.device\n        average_outputs = self.cumulative_average(\n          inputs, self.cumulative_average_mask(batch_size,\n                                               inputs_len).to(device).float()\n          if layer_cache is None else step, layer_cache=layer_cache)\n        average_outputs = self.average_layer(average_outputs)\n        gating_outputs = self.gating_layer(torch.cat((inputs,\n                                                      average_outputs), -1))\n        input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)\n        gating_outputs = torch.sigmoid(input_gate) * inputs + \\\n            torch.sigmoid(forget_gate) * average_outputs\n\n        return gating_outputs, average_outputs\n", "path": "onmt/modules/average_attn.py"}]}