We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Uninitialised bias parameters in RelPositionMultiHeadedAttention
## 🐛 Bug

The (u, v) bias terms in the MHA layer with relative position encodings are defined using `torch.Tensor`s:
https://github.com/facebookresearch/fairseq/blob/4fe8583396191c22011350248119db98ec1b5cb8/fairseq/modules/espnet_multihead_attention.py#L127-L128

Since `torch.Tensor` creates an un-initialised tensor, its use results in non-deterministic behaviour, poor initialisation, and nans if you have an unlucky init:
```python
>>>import torch
>>> torch.Tensor(1, 1).sum()
tensor(0.)
>>> torch.Tensor(10, 10).sum()
tensor(4.9312e-38)
>>> torch.Tensor(100, 100).sum()
tensor(nan)
```
Section 3.3 of the referenced paper [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/pdf/1901.02860.pdf) does not specify the initialisation for the bias terms. Since `torch.Tensor` is typically populated with zeros, setting the initialisation to `torch.zeros` should provide a fix that is close to the intended behaviour.


--- 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/espnet_multihead_attention.py`
Content:
```
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 # Copyright 2019 Shigeki Karita
5 #  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
6 
7 """Multi-Head Attention layer definition."""
8 
9 import math
10 import torch
11 from torch import nn
12 from fairseq.modules.rotary_positional_embedding import (
13     RotaryPositionalEmbedding,
14     apply_rotary_pos_emb,
15 )
16 
17 
18 class ESPNETMultiHeadedAttention(nn.Module):
19     """Multi-Head Attention layer.
20     Args:
21         n_head: The number of heads.
22         n_feat: The number of features.
23         dropout: Dropout rate.
24     """
25 
26     def __init__(self, n_feat, n_head, dropout):
27         """Construct an MultiHeadedAttention object."""
28         super(ESPNETMultiHeadedAttention, self).__init__()
29         assert n_feat % n_head == 0
30         # We assume d_v always equals d_k
31         self.d_k = n_feat // n_head
32         self.h = n_head
33         self.linear_q = nn.Linear(n_feat, n_feat)
34         self.linear_k = nn.Linear(n_feat, n_feat)
35         self.linear_v = nn.Linear(n_feat, n_feat)
36         self.linear_out = nn.Linear(n_feat, n_feat)
37         self.attn = None
38         self.dropout = nn.Dropout(p=dropout)
39 
40     def forward_qkv(self, query, key, value, **kwargs):
41         """Transform query, key and value.
42         Args:
43             query: Query tensor  B X T1 X C
44             key: Key tensor B X T2 X C
45             value: Value tensor  B X T2 X C
46         Returns:
47             torch.Tensor: Transformed query tensor  B X n_head X T1 X d_k
48             torch.Tensor: Transformed key tensor B X n_head X T2 X d_k
49             torch.Tensor: Transformed value tensor  B X n_head X T2 X d_k
50         """
51         n_batch = query.size(0)
52         q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
53         k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
54         v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
55         q = q.transpose(1, 2)  # (batch, head, time1, d_k)
56         k = k.transpose(1, 2)  # (batch, head, time2, d_k)
57         v = v.transpose(1, 2)  # (batch, head, time2, d_k)
58         return q, k, v
59 
60     def forward_attention(self, value, scores, mask):
61         """Compute attention context vector.
62         Args:
63             value: Transformed value B X n_head X T2 X d_k.
64             scores: Attention score  B X n_head X T1 X T2
65             mask: Mask  T2 X B
66         Returns:
67             torch.Tensor: Transformed value  B X T1 X d_model
68                 weighted by the attention score  B X T1 X T2
69         """
70         n_batch = value.size(0)
71         if mask is not None:
72             scores = scores.masked_fill(
73                 mask.unsqueeze(1).unsqueeze(2).to(bool),
74                 float("-inf"),  # (batch, head, time1, time2)
75             )
76             self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
77 
78         else:
79             self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
80         p_attn = self.dropout(self.attn)
81         x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
82         x = (
83             x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
84         )  # (batch, time1, d_model)
85 
86         return self.linear_out(x)  # (batch, time1, d_model)
87 
88     def forward(self, query, key, value, key_padding_mask=None, **kwargs):
89         """Compute scaled dot product attention.
90         Args:
91             query (torch.Tensor): Query tensor T X B X C
92             key (torch.Tensor): Key tensor T X B X C
93             value (torch.Tensor): Value tensor T X B X C
94             mask (torch.Tensor): Mask tensor T X B
95         Returns:
96             torch.Tensor: Output tensor T X B X D.
97         """
98         query = query.transpose(0, 1)
99         key = key.transpose(0, 1)
100         value = value.transpose(0, 1)
101 
102         q, k, v = self.forward_qkv(query, key, value)
103         scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
104         scores = self.forward_attention(v, scores, key_padding_mask)
105         scores = scores.transpose(0, 1)
106         return scores, None
107 
108 
109 class RelPositionMultiHeadedAttention(ESPNETMultiHeadedAttention):
110     """Multi-Head Attention layer with relative position encoding.
111     Paper: https://arxiv.org/abs/1901.02860
112     Args:
113         n_head: The number of heads.
114         n_feat: The number of features.
115         dropout: Dropout rate.
116         zero_triu: Whether to zero the upper triangular part of attention matrix.
117     """
118 
119     def __init__(self, n_feat, n_head, dropout, zero_triu=False):
120         """Construct an RelPositionMultiHeadedAttention object."""
121         super().__init__(n_feat, n_head, dropout)
122         self.zero_triu = zero_triu
123         # linear transformation for positional encoding
124         self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
125         # these two learnable bias are used in matrix c and matrix d
126         # as described in https://arxiv.org/abs/1901.02860 Section 3.3
127         self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
128         self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
129         torch.nn.init.xavier_uniform_(self.pos_bias_u)
130         torch.nn.init.xavier_uniform_(self.pos_bias_v)
131 
132     def rel_shift(self, x):
133         """Compute relative positional encoding.
134         Args:
135             x: Input tensor B X n_head X T X 2T-1
136         Returns:
137             torch.Tensor: Output tensor.
138         """
139         zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)
140         x_padded = torch.cat([zero_pad, x], dim=-1)
141 
142         x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
143         x = x_padded[:, :, 1:].view_as(x)[
144             :, :, :, : x.size(-1) // 2 + 1
145         ]  # only keep the positions from 0 to time2
146 
147         if self.zero_triu:
148             ones = torch.ones((x.size(2), x.size(3)), device=x.device)
149             x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]
150 
151         return x
152 
153     def forward(self, query, key, value, pos_emb, key_padding_mask=None, **kwargs):
154         """Compute scaled dot product attention.
155         Args:
156             query: Query tensor T X B X C
157             key: Key tensor T X B X C
158             value: Value tensor T X B X C
159             pos_emb: Positional embedding tensor B X 2T-1 X C
160             key_padding_mask: Mask tensor T X B
161         Returns:
162             torch.Tensor: Output tensor T X B X C.
163         """
164         query = query.transpose(0, 1)
165         key = key.transpose(0, 1)
166         value = value.transpose(0, 1)
167         pos_emb = pos_emb.transpose(0, 1)
168         q, k, v = self.forward_qkv(query, key, value)
169         q = q.transpose(1, 2)  # (batch, time1, head, d_k)
170         n_batch_pos = pos_emb.size(0)
171         p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
172         p = p.transpose(1, 2)  # (batch, head, 2*time1-1, d_k)
173 
174         # (batch, head, time1, d_k)
175         q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
176         # (batch, head, time1, d_k)
177         q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
178 
179         # compute attention score
180         # first compute matrix a and matrix c
181         # as described in https://arxiv.org/abs/1901.02860 Section 3.3
182         # (batch, head, time1, time2)
183         matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
184 
185         # compute matrix b and matrix d
186         # (batch, head, time1, 2*time1-1)
187         matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
188         matrix_bd = self.rel_shift(matrix_bd)
189 
190         scores = (matrix_ac + matrix_bd) / math.sqrt(
191             self.d_k
192         )  # (batch, head, time1, time2)
193 
194         scores = self.forward_attention(v, scores, key_padding_mask)
195         scores = scores.transpose(0, 1)
196         return scores, None
197 
198 
199 class RotaryPositionMultiHeadedAttention(ESPNETMultiHeadedAttention):
200     def __init__(
201         self,
202         n_feat,
203         n_head,
204         dropout,
205         precision,
206         rotary_emd_base=10000,
207     ):
208         """Construct an RotaryPositionMultiHeadedAttention object."""
209         super().__init__(n_feat, n_head, dropout)
210         precision = torch.float
211         self.rotary_ndims = self.d_k  # also try self.d_k//2
212         if precision == "fp16":
213             precision = torch.half
214 
215         self.rotary_emb = RotaryPositionalEmbedding(
216             self.rotary_ndims, base=rotary_emd_base, precision=precision
217         )
218 
219     def forward(self, query, key, value, key_padding_mask=None, **kwargs):
220         """Compute rotary position attention.
221         Args:
222             query: Query tensor T X B X C
223             key: Key tensor T X B X C
224             value: Value tensor T X B X C
225             key_padding_mask: Mask tensor T X B
226         Returns:
227             torch.Tensor: Output tensor T X B X D.
228         Notes:
229             Assumes self attn
230         """
231 
232         T, B, C = value.size()
233         query = query.view(T, B, self.h, self.d_k)
234         key = key.view(T, B, self.h, self.d_k)
235         value = value.view(T, B, self.h, self.d_k)
236         cos, sin = self.rotary_emb(value, seq_len=T)
237         query, key = apply_rotary_pos_emb(
238             query, key, cos, sin, offset=0
239         )  # offset is based on layer_past
240 
241         query = query.view(T, B, self.h * self.d_k)
242         key = key.view(T, B, self.h * self.d_k)
243         value = value.view(T, B, self.h * self.d_k)
244 
245         # TBD to BTD
246         query = query.transpose(0, 1)
247         key = key.transpose(0, 1)
248         value = value.transpose(0, 1)
249 
250         q, k, v = self.forward_qkv(query, key, value)
251         scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
252         scores = self.forward_attention(v, scores, key_padding_mask)
253         scores = scores.transpose(0, 1)
254         return scores, None
255
```
--- 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/espnet_multihead_attention.py b/fairseq/modules/espnet_multihead_attention.py
--- a/fairseq/modules/espnet_multihead_attention.py
+++ b/fairseq/modules/espnet_multihead_attention.py
@@ -124,8 +124,8 @@
         self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
         # these two learnable bias are used in matrix c and matrix d
         # as described in https://arxiv.org/abs/1901.02860 Section 3.3
-        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
-        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_u = nn.Parameter(torch.zeros(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.zeros(self.h, self.d_k))
         torch.nn.init.xavier_uniform_(self.pos_bias_u)
         torch.nn.init.xavier_uniform_(self.pos_bias_v)

{"golden_diff": "diff --git a/fairseq/modules/espnet_multihead_attention.py b/fairseq/modules/espnet_multihead_attention.py\n--- a/fairseq/modules/espnet_multihead_attention.py\n+++ b/fairseq/modules/espnet_multihead_attention.py\n@@ -124,8 +124,8 @@\n         self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)\n         # these two learnable bias are used in matrix c and matrix d\n         # as described in https://arxiv.org/abs/1901.02860 Section 3.3\n-        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))\n-        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))\n+        self.pos_bias_u = nn.Parameter(torch.zeros(self.h, self.d_k))\n+        self.pos_bias_v = nn.Parameter(torch.zeros(self.h, self.d_k))\n         torch.nn.init.xavier_uniform_(self.pos_bias_u)\n         torch.nn.init.xavier_uniform_(self.pos_bias_v)\n", "issue": "Uninitialised bias parameters in RelPositionMultiHeadedAttention\n## \ud83d\udc1b Bug\r\n\r\nThe (u, v) bias terms in the MHA layer with relative position encodings are defined using `torch.Tensor`s:\r\nhttps://github.com/facebookresearch/fairseq/blob/4fe8583396191c22011350248119db98ec1b5cb8/fairseq/modules/espnet_multihead_attention.py#L127-L128\r\n\r\nSince `torch.Tensor` creates an un-initialised tensor, its use results in non-deterministic behaviour, poor initialisation, and nans if you have an unlucky init:\r\n```python\r\n>>>import torch\r\n>>> torch.Tensor(1, 1).sum()\r\ntensor(0.)\r\n>>> torch.Tensor(10, 10).sum()\r\ntensor(4.9312e-38)\r\n>>> torch.Tensor(100, 100).sum()\r\ntensor(nan)\r\n```\r\nSection 3.3 of the referenced paper [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/pdf/1901.02860.pdf) does not specify the initialisation for the bias terms. Since `torch.Tensor` is typically populated with zeros, setting the initialisation to `torch.zeros` should provide a fix that is close to the intended behaviour.\r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\n# Copyright 2019 Shigeki Karita\n#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)\n\n\"\"\"Multi-Head Attention layer definition.\"\"\"\n\nimport math\nimport torch\nfrom torch import nn\nfrom fairseq.modules.rotary_positional_embedding import (\n    RotaryPositionalEmbedding,\n    apply_rotary_pos_emb,\n)\n\n\nclass ESPNETMultiHeadedAttention(nn.Module):\n    \"\"\"Multi-Head Attention layer.\n    Args:\n        n_head: The number of heads.\n        n_feat: The number of features.\n        dropout: Dropout rate.\n    \"\"\"\n\n    def __init__(self, n_feat, n_head, dropout):\n        \"\"\"Construct an MultiHeadedAttention object.\"\"\"\n        super(ESPNETMultiHeadedAttention, self).__init__()\n        assert n_feat % n_head == 0\n        # We assume d_v always equals d_k\n        self.d_k = n_feat // n_head\n        self.h = n_head\n        self.linear_q = nn.Linear(n_feat, n_feat)\n        self.linear_k = nn.Linear(n_feat, n_feat)\n        self.linear_v = nn.Linear(n_feat, n_feat)\n        self.linear_out = nn.Linear(n_feat, n_feat)\n        self.attn = None\n        self.dropout = nn.Dropout(p=dropout)\n\n    def forward_qkv(self, query, key, value, **kwargs):\n        \"\"\"Transform query, key and value.\n        Args:\n            query: Query tensor  B X T1 X C\n            key: Key tensor B X T2 X C\n            value: Value tensor  B X T2 X C\n        Returns:\n            torch.Tensor: Transformed query tensor  B X n_head X T1 X d_k\n            torch.Tensor: Transformed key tensor B X n_head X T2 X d_k\n            torch.Tensor: Transformed value tensor  B X n_head X T2 X d_k\n        \"\"\"\n        n_batch = query.size(0)\n        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)\n        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)\n        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)\n        q = q.transpose(1, 2)  # (batch, head, time1, d_k)\n        k = k.transpose(1, 2)  # (batch, head, time2, d_k)\n        v = v.transpose(1, 2)  # (batch, head, time2, d_k)\n        return q, k, v\n\n    def forward_attention(self, value, scores, mask):\n        \"\"\"Compute attention context vector.\n        Args:\n            value: Transformed value B X n_head X T2 X d_k.\n            scores: Attention score  B X n_head X T1 X T2\n            mask: Mask  T2 X B\n        Returns:\n            torch.Tensor: Transformed value  B X T1 X d_model\n                weighted by the attention score  B X T1 X T2\n        \"\"\"\n        n_batch = value.size(0)\n        if mask is not None:\n            scores = scores.masked_fill(\n                mask.unsqueeze(1).unsqueeze(2).to(bool),\n                float(\"-inf\"),  # (batch, head, time1, time2)\n            )\n            self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)\n\n        else:\n            self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)\n        p_attn = self.dropout(self.attn)\n        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)\n        x = (\n            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)\n        )  # (batch, time1, d_model)\n\n        return self.linear_out(x)  # (batch, time1, d_model)\n\n    def forward(self, query, key, value, key_padding_mask=None, **kwargs):\n        \"\"\"Compute scaled dot product attention.\n        Args:\n            query (torch.Tensor): Query tensor T X B X C\n            key (torch.Tensor): Key tensor T X B X C\n            value (torch.Tensor): Value tensor T X B X C\n            mask (torch.Tensor): Mask tensor T X B\n        Returns:\n            torch.Tensor: Output tensor T X B X D.\n        \"\"\"\n        query = query.transpose(0, 1)\n        key = key.transpose(0, 1)\n        value = value.transpose(0, 1)\n\n        q, k, v = self.forward_qkv(query, key, value)\n        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n        scores = self.forward_attention(v, scores, key_padding_mask)\n        scores = scores.transpose(0, 1)\n        return scores, None\n\n\nclass RelPositionMultiHeadedAttention(ESPNETMultiHeadedAttention):\n    \"\"\"Multi-Head Attention layer with relative position encoding.\n    Paper: https://arxiv.org/abs/1901.02860\n    Args:\n        n_head: The number of heads.\n        n_feat: The number of features.\n        dropout: Dropout rate.\n        zero_triu: Whether to zero the upper triangular part of attention matrix.\n    \"\"\"\n\n    def __init__(self, n_feat, n_head, dropout, zero_triu=False):\n        \"\"\"Construct an RelPositionMultiHeadedAttention object.\"\"\"\n        super().__init__(n_feat, n_head, dropout)\n        self.zero_triu = zero_triu\n        # linear transformation for positional encoding\n        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)\n        # these two learnable bias are used in matrix c and matrix d\n        # as described in https://arxiv.org/abs/1901.02860 Section 3.3\n        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))\n        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))\n        torch.nn.init.xavier_uniform_(self.pos_bias_u)\n        torch.nn.init.xavier_uniform_(self.pos_bias_v)\n\n    def rel_shift(self, x):\n        \"\"\"Compute relative positional encoding.\n        Args:\n            x: Input tensor B X n_head X T X 2T-1\n        Returns:\n            torch.Tensor: Output tensor.\n        \"\"\"\n        zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)\n        x_padded = torch.cat([zero_pad, x], dim=-1)\n\n        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))\n        x = x_padded[:, :, 1:].view_as(x)[\n            :, :, :, : x.size(-1) // 2 + 1\n        ]  # only keep the positions from 0 to time2\n\n        if self.zero_triu:\n            ones = torch.ones((x.size(2), x.size(3)), device=x.device)\n            x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]\n\n        return x\n\n    def forward(self, query, key, value, pos_emb, key_padding_mask=None, **kwargs):\n        \"\"\"Compute scaled dot product attention.\n        Args:\n            query: Query tensor T X B X C\n            key: Key tensor T X B X C\n            value: Value tensor T X B X C\n            pos_emb: Positional embedding tensor B X 2T-1 X C\n            key_padding_mask: Mask tensor T X B\n        Returns:\n            torch.Tensor: Output tensor T X B X C.\n        \"\"\"\n        query = query.transpose(0, 1)\n        key = key.transpose(0, 1)\n        value = value.transpose(0, 1)\n        pos_emb = pos_emb.transpose(0, 1)\n        q, k, v = self.forward_qkv(query, key, value)\n        q = q.transpose(1, 2)  # (batch, time1, head, d_k)\n        n_batch_pos = pos_emb.size(0)\n        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)\n        p = p.transpose(1, 2)  # (batch, head, 2*time1-1, d_k)\n\n        # (batch, head, time1, d_k)\n        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)\n        # (batch, head, time1, d_k)\n        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)\n\n        # compute attention score\n        # first compute matrix a and matrix c\n        # as described in https://arxiv.org/abs/1901.02860 Section 3.3\n        # (batch, head, time1, time2)\n        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))\n\n        # compute matrix b and matrix d\n        # (batch, head, time1, 2*time1-1)\n        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))\n        matrix_bd = self.rel_shift(matrix_bd)\n\n        scores = (matrix_ac + matrix_bd) / math.sqrt(\n            self.d_k\n        )  # (batch, head, time1, time2)\n\n        scores = self.forward_attention(v, scores, key_padding_mask)\n        scores = scores.transpose(0, 1)\n        return scores, None\n\n\nclass RotaryPositionMultiHeadedAttention(ESPNETMultiHeadedAttention):\n    def __init__(\n        self,\n        n_feat,\n        n_head,\n        dropout,\n        precision,\n        rotary_emd_base=10000,\n    ):\n        \"\"\"Construct an RotaryPositionMultiHeadedAttention object.\"\"\"\n        super().__init__(n_feat, n_head, dropout)\n        precision = torch.float\n        self.rotary_ndims = self.d_k  # also try self.d_k//2\n        if precision == \"fp16\":\n            precision = torch.half\n\n        self.rotary_emb = RotaryPositionalEmbedding(\n            self.rotary_ndims, base=rotary_emd_base, precision=precision\n        )\n\n    def forward(self, query, key, value, key_padding_mask=None, **kwargs):\n        \"\"\"Compute rotary position attention.\n        Args:\n            query: Query tensor T X B X C\n            key: Key tensor T X B X C\n            value: Value tensor T X B X C\n            key_padding_mask: Mask tensor T X B\n        Returns:\n            torch.Tensor: Output tensor T X B X D.\n        Notes:\n            Assumes self attn\n        \"\"\"\n\n        T, B, C = value.size()\n        query = query.view(T, B, self.h, self.d_k)\n        key = key.view(T, B, self.h, self.d_k)\n        value = value.view(T, B, self.h, self.d_k)\n        cos, sin = self.rotary_emb(value, seq_len=T)\n        query, key = apply_rotary_pos_emb(\n            query, key, cos, sin, offset=0\n        )  # offset is based on layer_past\n\n        query = query.view(T, B, self.h * self.d_k)\n        key = key.view(T, B, self.h * self.d_k)\n        value = value.view(T, B, self.h * self.d_k)\n\n        # TBD to BTD\n        query = query.transpose(0, 1)\n        key = key.transpose(0, 1)\n        value = value.transpose(0, 1)\n\n        q, k, v = self.forward_qkv(query, key, value)\n        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n        scores = self.forward_attention(v, scores, key_padding_mask)\n        scores = scores.transpose(0, 1)\n        return scores, None\n", "path": "fairseq/modules/espnet_multihead_attention.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\n# Copyright 2019 Shigeki Karita\n#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)\n\n\"\"\"Multi-Head Attention layer definition.\"\"\"\n\nimport math\nimport torch\nfrom torch import nn\nfrom fairseq.modules.rotary_positional_embedding import (\n    RotaryPositionalEmbedding,\n    apply_rotary_pos_emb,\n)\n\n\nclass ESPNETMultiHeadedAttention(nn.Module):\n    \"\"\"Multi-Head Attention layer.\n    Args:\n        n_head: The number of heads.\n        n_feat: The number of features.\n        dropout: Dropout rate.\n    \"\"\"\n\n    def __init__(self, n_feat, n_head, dropout):\n        \"\"\"Construct an MultiHeadedAttention object.\"\"\"\n        super(ESPNETMultiHeadedAttention, self).__init__()\n        assert n_feat % n_head == 0\n        # We assume d_v always equals d_k\n        self.d_k = n_feat // n_head\n        self.h = n_head\n        self.linear_q = nn.Linear(n_feat, n_feat)\n        self.linear_k = nn.Linear(n_feat, n_feat)\n        self.linear_v = nn.Linear(n_feat, n_feat)\n        self.linear_out = nn.Linear(n_feat, n_feat)\n        self.attn = None\n        self.dropout = nn.Dropout(p=dropout)\n\n    def forward_qkv(self, query, key, value, **kwargs):\n        \"\"\"Transform query, key and value.\n        Args:\n            query: Query tensor  B X T1 X C\n            key: Key tensor B X T2 X C\n            value: Value tensor  B X T2 X C\n        Returns:\n            torch.Tensor: Transformed query tensor  B X n_head X T1 X d_k\n            torch.Tensor: Transformed key tensor B X n_head X T2 X d_k\n            torch.Tensor: Transformed value tensor  B X n_head X T2 X d_k\n        \"\"\"\n        n_batch = query.size(0)\n        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)\n        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)\n        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)\n        q = q.transpose(1, 2)  # (batch, head, time1, d_k)\n        k = k.transpose(1, 2)  # (batch, head, time2, d_k)\n        v = v.transpose(1, 2)  # (batch, head, time2, d_k)\n        return q, k, v\n\n    def forward_attention(self, value, scores, mask):\n        \"\"\"Compute attention context vector.\n        Args:\n            value: Transformed value B X n_head X T2 X d_k.\n            scores: Attention score  B X n_head X T1 X T2\n            mask: Mask  T2 X B\n        Returns:\n            torch.Tensor: Transformed value  B X T1 X d_model\n                weighted by the attention score  B X T1 X T2\n        \"\"\"\n        n_batch = value.size(0)\n        if mask is not None:\n            scores = scores.masked_fill(\n                mask.unsqueeze(1).unsqueeze(2).to(bool),\n                float(\"-inf\"),  # (batch, head, time1, time2)\n            )\n            self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)\n\n        else:\n            self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)\n        p_attn = self.dropout(self.attn)\n        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)\n        x = (\n            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)\n        )  # (batch, time1, d_model)\n\n        return self.linear_out(x)  # (batch, time1, d_model)\n\n    def forward(self, query, key, value, key_padding_mask=None, **kwargs):\n        \"\"\"Compute scaled dot product attention.\n        Args:\n            query (torch.Tensor): Query tensor T X B X C\n            key (torch.Tensor): Key tensor T X B X C\n            value (torch.Tensor): Value tensor T X B X C\n            mask (torch.Tensor): Mask tensor T X B\n        Returns:\n            torch.Tensor: Output tensor T X B X D.\n        \"\"\"\n        query = query.transpose(0, 1)\n        key = key.transpose(0, 1)\n        value = value.transpose(0, 1)\n\n        q, k, v = self.forward_qkv(query, key, value)\n        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n        scores = self.forward_attention(v, scores, key_padding_mask)\n        scores = scores.transpose(0, 1)\n        return scores, None\n\n\nclass RelPositionMultiHeadedAttention(ESPNETMultiHeadedAttention):\n    \"\"\"Multi-Head Attention layer with relative position encoding.\n    Paper: https://arxiv.org/abs/1901.02860\n    Args:\n        n_head: The number of heads.\n        n_feat: The number of features.\n        dropout: Dropout rate.\n        zero_triu: Whether to zero the upper triangular part of attention matrix.\n    \"\"\"\n\n    def __init__(self, n_feat, n_head, dropout, zero_triu=False):\n        \"\"\"Construct an RelPositionMultiHeadedAttention object.\"\"\"\n        super().__init__(n_feat, n_head, dropout)\n        self.zero_triu = zero_triu\n        # linear transformation for positional encoding\n        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)\n        # these two learnable bias are used in matrix c and matrix d\n        # as described in https://arxiv.org/abs/1901.02860 Section 3.3\n        self.pos_bias_u = nn.Parameter(torch.zeros(self.h, self.d_k))\n        self.pos_bias_v = nn.Parameter(torch.zeros(self.h, self.d_k))\n        torch.nn.init.xavier_uniform_(self.pos_bias_u)\n        torch.nn.init.xavier_uniform_(self.pos_bias_v)\n\n    def rel_shift(self, x):\n        \"\"\"Compute relative positional encoding.\n        Args:\n            x: Input tensor B X n_head X T X 2T-1\n        Returns:\n            torch.Tensor: Output tensor.\n        \"\"\"\n        zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)\n        x_padded = torch.cat([zero_pad, x], dim=-1)\n\n        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))\n        x = x_padded[:, :, 1:].view_as(x)[\n            :, :, :, : x.size(-1) // 2 + 1\n        ]  # only keep the positions from 0 to time2\n\n        if self.zero_triu:\n            ones = torch.ones((x.size(2), x.size(3)), device=x.device)\n            x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]\n\n        return x\n\n    def forward(self, query, key, value, pos_emb, key_padding_mask=None, **kwargs):\n        \"\"\"Compute scaled dot product attention.\n        Args:\n            query: Query tensor T X B X C\n            key: Key tensor T X B X C\n            value: Value tensor T X B X C\n            pos_emb: Positional embedding tensor B X 2T-1 X C\n            key_padding_mask: Mask tensor T X B\n        Returns:\n            torch.Tensor: Output tensor T X B X C.\n        \"\"\"\n        query = query.transpose(0, 1)\n        key = key.transpose(0, 1)\n        value = value.transpose(0, 1)\n        pos_emb = pos_emb.transpose(0, 1)\n        q, k, v = self.forward_qkv(query, key, value)\n        q = q.transpose(1, 2)  # (batch, time1, head, d_k)\n        n_batch_pos = pos_emb.size(0)\n        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)\n        p = p.transpose(1, 2)  # (batch, head, 2*time1-1, d_k)\n\n        # (batch, head, time1, d_k)\n        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)\n        # (batch, head, time1, d_k)\n        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)\n\n        # compute attention score\n        # first compute matrix a and matrix c\n        # as described in https://arxiv.org/abs/1901.02860 Section 3.3\n        # (batch, head, time1, time2)\n        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))\n\n        # compute matrix b and matrix d\n        # (batch, head, time1, 2*time1-1)\n        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))\n        matrix_bd = self.rel_shift(matrix_bd)\n\n        scores = (matrix_ac + matrix_bd) / math.sqrt(\n            self.d_k\n        )  # (batch, head, time1, time2)\n\n        scores = self.forward_attention(v, scores, key_padding_mask)\n        scores = scores.transpose(0, 1)\n        return scores, None\n\n\nclass RotaryPositionMultiHeadedAttention(ESPNETMultiHeadedAttention):\n    def __init__(\n        self,\n        n_feat,\n        n_head,\n        dropout,\n        precision,\n        rotary_emd_base=10000,\n    ):\n        \"\"\"Construct an RotaryPositionMultiHeadedAttention object.\"\"\"\n        super().__init__(n_feat, n_head, dropout)\n        precision = torch.float\n        self.rotary_ndims = self.d_k  # also try self.d_k//2\n        if precision == \"fp16\":\n            precision = torch.half\n\n        self.rotary_emb = RotaryPositionalEmbedding(\n            self.rotary_ndims, base=rotary_emd_base, precision=precision\n        )\n\n    def forward(self, query, key, value, key_padding_mask=None, **kwargs):\n        \"\"\"Compute rotary position attention.\n        Args:\n            query: Query tensor T X B X C\n            key: Key tensor T X B X C\n            value: Value tensor T X B X C\n            key_padding_mask: Mask tensor T X B\n        Returns:\n            torch.Tensor: Output tensor T X B X D.\n        Notes:\n            Assumes self attn\n        \"\"\"\n\n        T, B, C = value.size()\n        query = query.view(T, B, self.h, self.d_k)\n        key = key.view(T, B, self.h, self.d_k)\n        value = value.view(T, B, self.h, self.d_k)\n        cos, sin = self.rotary_emb(value, seq_len=T)\n        query, key = apply_rotary_pos_emb(\n            query, key, cos, sin, offset=0\n        )  # offset is based on layer_past\n\n        query = query.view(T, B, self.h * self.d_k)\n        key = key.view(T, B, self.h * self.d_k)\n        value = value.view(T, B, self.h * self.d_k)\n\n        # TBD to BTD\n        query = query.transpose(0, 1)\n        key = key.transpose(0, 1)\n        value = value.transpose(0, 1)\n\n        q, k, v = self.forward_qkv(query, key, value)\n        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n        scores = self.forward_attention(v, scores, key_padding_mask)\n        scores = scores.transpose(0, 1)\n        return scores, None\n", "path": "fairseq/modules/espnet_multihead_attention.py"}]}