We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[fairseq] Guard call to `shape_as_tensor` with `is_in_onnx_export()`
This is a no-op in eager and in ONNX export, but it's better for other
tracers if this is preserved as shapes directly instead of converted to
a tensor.

There is a little annoying code duplication with
`torch.jit.is_scripting()`, which is unforunately necessary because we
didn't implement compile-time short circuiting correctly in TorchScript
lol.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `fairseq/modules/sinusoidal_positional_embedding.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates.
2 #
3 # This source code is licensed under the MIT license found in the
4 # LICENSE file in the root directory of this source tree.
5 
6 import math
7 from typing import Any, Optional
8 
9 import torch
10 import torch.onnx.operators
11 from fairseq import utils
12 from torch import Tensor, nn
13 
14 
15 class SinusoidalPositionalEmbedding(nn.Module):
16     """This module produces sinusoidal positional embeddings of any length.
17 
18     Padding symbols are ignored.
19     """
20 
21     def __init__(self, embedding_dim, padding_idx, init_size=1024):
22         super().__init__()
23         self.embedding_dim = embedding_dim
24         self.padding_idx = padding_idx if padding_idx is not None else 0
25         self.weights = SinusoidalPositionalEmbedding.get_embedding(
26             init_size, embedding_dim, padding_idx
27         )
28         self.onnx_trace = False
29         self.register_buffer("_float_tensor", torch.FloatTensor(1))
30         self.max_positions = int(1e5)
31 
32     def prepare_for_onnx_export_(self):
33         self.onnx_trace = True
34 
35     @staticmethod
36     def get_embedding(
37         num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None
38     ):
39         """Build sinusoidal embeddings.
40 
41         This matches the implementation in tensor2tensor, but differs slightly
42         from the description in Section 3.5 of "Attention Is All You Need".
43         """
44         half_dim = embedding_dim // 2
45         emb = math.log(10000) / (half_dim - 1)
46         emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
47         emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(
48             1
49         ) * emb.unsqueeze(0)
50         emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(
51             num_embeddings, -1
52         )
53         if embedding_dim % 2 == 1:
54             # zero pad
55             emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
56         if padding_idx is not None:
57             emb[padding_idx, :] = 0
58         return emb
59 
60     def forward(
61         self,
62         input,
63         incremental_state: Optional[Any] = None,
64         timestep: Optional[Tensor] = None,
65         positions: Optional[Any] = None,
66     ):
67         """Input is expected to be of size [bsz x seqlen]."""
68         if torch.jit.is_scripting():
69             bspair = torch.onnx.operators.shape_as_tensor(input)
70         elif torch.onnx.is_in_onnx_export():
71             bspair = torch.onnx.operators.shape_as_tensor(input)
72         else:
73             bspair = input.size()
74         bsz, seq_len = bspair[0], bspair[1]
75         max_pos = self.padding_idx + 1 + seq_len
76         if self.weights is None or max_pos > self.weights.size(0):
77             # recompute/expand embeddings if needed
78             self.weights = SinusoidalPositionalEmbedding.get_embedding(
79                 max_pos, self.embedding_dim, self.padding_idx
80             )
81         self.weights = self.weights.to(self._float_tensor)
82 
83         if incremental_state is not None:
84             # positions is the same for every token when decoding a single step
85             pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
86             if self.onnx_trace:
87                 return (
88                     self.weights.index_select(index=self.padding_idx + pos, dim=0)
89                     .unsqueeze(1)
90                     .repeat(bsz, 1, 1)
91                 )
92             return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)
93 
94         positions = utils.make_positions(
95             input, self.padding_idx, onnx_trace=self.onnx_trace
96         )
97         if self.onnx_trace:
98             flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))
99             embedding_shape = torch.cat(
100                 (bsz.view(1), seq_len.view(1), torch.tensor([-1], dtype=torch.long))
101             )
102             embeddings = torch.onnx.operators.reshape_from_tensor_shape(
103                 flat_embeddings, embedding_shape
104             )
105             return embeddings
106         return (
107             self.weights.index_select(0, positions.view(-1))
108             .view(bsz, seq_len, -1)
109             .detach()
110         )
111
```
--- 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()

```

diff --git a/fairseq/modules/sinusoidal_positional_embedding.py b/fairseq/modules/sinusoidal_positional_embedding.py
--- a/fairseq/modules/sinusoidal_positional_embedding.py
+++ b/fairseq/modules/sinusoidal_positional_embedding.py
@@ -97,7 +97,7 @@
         if self.onnx_trace:
             flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))
             embedding_shape = torch.cat(
-                (bsz.view(1), seq_len.view(1), torch.tensor([-1], dtype=torch.long))
+                (bsz, seq_len, torch.tensor([-1], dtype=torch.long))
             )
             embeddings = torch.onnx.operators.reshape_from_tensor_shape(
                 flat_embeddings, embedding_shape

{"golden_diff": "diff --git a/fairseq/modules/sinusoidal_positional_embedding.py b/fairseq/modules/sinusoidal_positional_embedding.py\n--- a/fairseq/modules/sinusoidal_positional_embedding.py\n+++ b/fairseq/modules/sinusoidal_positional_embedding.py\n@@ -97,7 +97,7 @@\n         if self.onnx_trace:\n             flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))\n             embedding_shape = torch.cat(\n-                (bsz.view(1), seq_len.view(1), torch.tensor([-1], dtype=torch.long))\n+                (bsz, seq_len, torch.tensor([-1], dtype=torch.long))\n             )\n             embeddings = torch.onnx.operators.reshape_from_tensor_shape(\n                 flat_embeddings, embedding_shape\n", "issue": "[fairseq] Guard call to `shape_as_tensor` with `is_in_onnx_export()`\nThis is a no-op in eager and in ONNX export, but it's better for other\ntracers if this is preserved as shapes directly instead of converted to\na tensor.\n\nThere is a little annoying code duplication with\n`torch.jit.is_scripting()`, which is unforunately necessary because we\ndidn't implement compile-time short circuiting correctly in TorchScript\nlol.\n\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom typing import Any, Optional\n\nimport torch\nimport torch.onnx.operators\nfrom fairseq import utils\nfrom torch import Tensor, nn\n\n\nclass SinusoidalPositionalEmbedding(nn.Module):\n    \"\"\"This module produces sinusoidal positional embeddings of any length.\n\n    Padding symbols are ignored.\n    \"\"\"\n\n    def __init__(self, embedding_dim, padding_idx, init_size=1024):\n        super().__init__()\n        self.embedding_dim = embedding_dim\n        self.padding_idx = padding_idx if padding_idx is not None else 0\n        self.weights = SinusoidalPositionalEmbedding.get_embedding(\n            init_size, embedding_dim, padding_idx\n        )\n        self.onnx_trace = False\n        self.register_buffer(\"_float_tensor\", torch.FloatTensor(1))\n        self.max_positions = int(1e5)\n\n    def prepare_for_onnx_export_(self):\n        self.onnx_trace = True\n\n    @staticmethod\n    def get_embedding(\n        num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None\n    ):\n        \"\"\"Build sinusoidal embeddings.\n\n        This matches the implementation in tensor2tensor, but differs slightly\n        from the description in Section 3.5 of \"Attention Is All You Need\".\n        \"\"\"\n        half_dim = embedding_dim // 2\n        emb = math.log(10000) / (half_dim - 1)\n        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)\n        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(\n            1\n        ) * emb.unsqueeze(0)\n        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(\n            num_embeddings, -1\n        )\n        if embedding_dim % 2 == 1:\n            # zero pad\n            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)\n        if padding_idx is not None:\n            emb[padding_idx, :] = 0\n        return emb\n\n    def forward(\n        self,\n        input,\n        incremental_state: Optional[Any] = None,\n        timestep: Optional[Tensor] = None,\n        positions: Optional[Any] = None,\n    ):\n        \"\"\"Input is expected to be of size [bsz x seqlen].\"\"\"\n        if torch.jit.is_scripting():\n            bspair = torch.onnx.operators.shape_as_tensor(input)\n        elif torch.onnx.is_in_onnx_export():\n            bspair = torch.onnx.operators.shape_as_tensor(input)\n        else:\n            bspair = input.size()\n        bsz, seq_len = bspair[0], bspair[1]\n        max_pos = self.padding_idx + 1 + seq_len\n        if self.weights is None or max_pos > self.weights.size(0):\n            # recompute/expand embeddings if needed\n            self.weights = SinusoidalPositionalEmbedding.get_embedding(\n                max_pos, self.embedding_dim, self.padding_idx\n            )\n        self.weights = self.weights.to(self._float_tensor)\n\n        if incremental_state is not None:\n            # positions is the same for every token when decoding a single step\n            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len\n            if self.onnx_trace:\n                return (\n                    self.weights.index_select(index=self.padding_idx + pos, dim=0)\n                    .unsqueeze(1)\n                    .repeat(bsz, 1, 1)\n                )\n            return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)\n\n        positions = utils.make_positions(\n            input, self.padding_idx, onnx_trace=self.onnx_trace\n        )\n        if self.onnx_trace:\n            flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))\n            embedding_shape = torch.cat(\n                (bsz.view(1), seq_len.view(1), torch.tensor([-1], dtype=torch.long))\n            )\n            embeddings = torch.onnx.operators.reshape_from_tensor_shape(\n                flat_embeddings, embedding_shape\n            )\n            return embeddings\n        return (\n            self.weights.index_select(0, positions.view(-1))\n            .view(bsz, seq_len, -1)\n            .detach()\n        )\n", "path": "fairseq/modules/sinusoidal_positional_embedding.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom typing import Any, Optional\n\nimport torch\nimport torch.onnx.operators\nfrom fairseq import utils\nfrom torch import Tensor, nn\n\n\nclass SinusoidalPositionalEmbedding(nn.Module):\n    \"\"\"This module produces sinusoidal positional embeddings of any length.\n\n    Padding symbols are ignored.\n    \"\"\"\n\n    def __init__(self, embedding_dim, padding_idx, init_size=1024):\n        super().__init__()\n        self.embedding_dim = embedding_dim\n        self.padding_idx = padding_idx if padding_idx is not None else 0\n        self.weights = SinusoidalPositionalEmbedding.get_embedding(\n            init_size, embedding_dim, padding_idx\n        )\n        self.onnx_trace = False\n        self.register_buffer(\"_float_tensor\", torch.FloatTensor(1))\n        self.max_positions = int(1e5)\n\n    def prepare_for_onnx_export_(self):\n        self.onnx_trace = True\n\n    @staticmethod\n    def get_embedding(\n        num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None\n    ):\n        \"\"\"Build sinusoidal embeddings.\n\n        This matches the implementation in tensor2tensor, but differs slightly\n        from the description in Section 3.5 of \"Attention Is All You Need\".\n        \"\"\"\n        half_dim = embedding_dim // 2\n        emb = math.log(10000) / (half_dim - 1)\n        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)\n        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(\n            1\n        ) * emb.unsqueeze(0)\n        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(\n            num_embeddings, -1\n        )\n        if embedding_dim % 2 == 1:\n            # zero pad\n            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)\n        if padding_idx is not None:\n            emb[padding_idx, :] = 0\n        return emb\n\n    def forward(\n        self,\n        input,\n        incremental_state: Optional[Any] = None,\n        timestep: Optional[Tensor] = None,\n        positions: Optional[Any] = None,\n    ):\n        \"\"\"Input is expected to be of size [bsz x seqlen].\"\"\"\n        if torch.jit.is_scripting():\n            bspair = torch.onnx.operators.shape_as_tensor(input)\n        elif torch.onnx.is_in_onnx_export():\n            bspair = torch.onnx.operators.shape_as_tensor(input)\n        else:\n            bspair = input.size()\n        bsz, seq_len = bspair[0], bspair[1]\n        max_pos = self.padding_idx + 1 + seq_len\n        if self.weights is None or max_pos > self.weights.size(0):\n            # recompute/expand embeddings if needed\n            self.weights = SinusoidalPositionalEmbedding.get_embedding(\n                max_pos, self.embedding_dim, self.padding_idx\n            )\n        self.weights = self.weights.to(self._float_tensor)\n\n        if incremental_state is not None:\n            # positions is the same for every token when decoding a single step\n            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len\n            if self.onnx_trace:\n                return (\n                    self.weights.index_select(index=self.padding_idx + pos, dim=0)\n                    .unsqueeze(1)\n                    .repeat(bsz, 1, 1)\n                )\n            return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)\n\n        positions = utils.make_positions(\n            input, self.padding_idx, onnx_trace=self.onnx_trace\n        )\n        if self.onnx_trace:\n            flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))\n            embedding_shape = torch.cat(\n                (bsz, seq_len, torch.tensor([-1], dtype=torch.long))\n            )\n            embeddings = torch.onnx.operators.reshape_from_tensor_shape(\n                flat_embeddings, embedding_shape\n            )\n            return embeddings\n        return (\n            self.weights.index_select(0, positions.view(-1))\n            .view(bsz, seq_len, -1)\n            .detach()\n        )\n", "path": "fairseq/modules/sinusoidal_positional_embedding.py"}]}