We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Cannot pickle nystrom based model
When I `torch.save` my model with `nystrom` attention I get the following pickle error:

```py
AttributeError: Can't pickle local object 'get_avg_pool.<locals>.avg_pool'
```

I believe coming from this function:

https://github.com/facebookresearch/xformers/blob/9232b2d27a775a43173e2fd86c03251ab64f7ede/xformers/components/attention/nystrom.py#L60

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `xformers/components/attention/nystrom.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
2 #
3 # This source code is licensed under the BSD license found in the
4 # LICENSE file in the root directory of this source tree.
5 
6 
7 import logging
8 from dataclasses import dataclass
9 from typing import Optional
10 
11 import torch
12 import torch.nn as nn
13 
14 from xformers.components.attention import Attention, AttentionConfig, register_attention
15 from xformers.components.attention.core import (
16     scaled_dot_product_attention,
17     scaled_query_key_softmax,
18 )
19 from xformers.components.attention.utils import (
20     bool_mask_to_additive,
21     iterative_pinv,
22     reshape_key_padding_mask,
23 )
24 
25 
26 @dataclass
27 class NystromSelfAttentionConfig(AttentionConfig):
28     """
29     num_heads               Number of heads.
30     num_landmarks           Number of landmarks to use for softmax approximation. 64 often sufficient for a good
31                             approximation according to https://arxiv.org/pdf/2102.03902.pdf.
32     causal                  Apply a causal mask, in that the attention cannot be applied to the future.
33     use_razavi_pinverse     If true, use iterative method from (Razavi et al. 2014) to approximate the Moore-Penrose
34                             inverse, otherwise use standard torch inverse.
35     pinverse_original_init  True if using original initialization when calculating Moore-Penrose pseudo inverse using
36                             method from (Razavi et al. 2014).
37                             False if using exact coefficient computation (leads to faster convergence).
38     inv_iterations          Number of iterations for calculating the Moore-Penrose pseudo inverse.
39     v_skip_connection       A module that will take V as input and will be added as a skip connection to the
40                             softmax approximation. A skip connection is added in the paper to help with training.
41     conv_kernel_size        Kernel size for convolution optionally added to help in training.
42                             If v_skip_connection is not specified, this will be used to define the default
43                             depth wise convolution used as a skip connection.
44                             If both conv_kernel_size and v_skip_connection are None, no skip connection will
45                             be added.
46     landmark_pooling        Which module to use when computing landmarks. Default is AdaptiveAvgPool2d.
47     """
48 
49     num_heads: int
50     num_landmarks: Optional[int]
51     landmark_pooling: Optional[nn.Module]
52     causal: Optional[bool]
53     pinverse_original_init: Optional[bool]
54     inv_iterations: Optional[int]
55     v_skip_connection: Optional[nn.Module]
56     conv_kernel_size: Optional[int]
57     use_razavi_pinverse: Optional[bool]
58 
59 
60 def get_avg_pool(n: int):
61     def avg_pool(x: torch.Tensor):
62         # Average independently for every segment in the sequence dimension
63         seq_len = x.shape[1]
64         head_dim = x.shape[2]
65         segments = seq_len // n
66 
67         # Dimensions are a match
68         if seq_len % n == 0:
69             return x.reshape(
70                 -1,
71                 n,
72                 segments,
73                 head_dim,
74             ).mean(dim=-2)
75 
76         # Handle the last segment boundary being off
77         n_round = n - seq_len % n
78 
79         x_avg_round = (
80             x[:, : n_round * segments, :]
81             .reshape(-1, n_round, segments, head_dim)
82             .mean(dim=-2)
83         )
84         x_avg_off = (
85             x[:, n_round * segments :, :]
86             .reshape(-1, n - n_round, segments + 1, head_dim)
87             .mean(dim=-2)
88         )
89         return torch.cat((x_avg_round, x_avg_off), dim=-2)
90 
91     return avg_pool
92 
93 
94 @register_attention("nystrom", NystromSelfAttentionConfig)
95 class NystromAttention(Attention):
96     # TODO: update defaults for use_razavi_pinverse and inv_iterations
97     def __init__(
98         self,
99         dropout: float,
100         num_heads: int,
101         num_landmarks: int = 64,
102         landmark_pooling: Optional[nn.Module] = None,
103         causal: bool = False,
104         use_razavi_pinverse: bool = True,
105         pinverse_original_init: bool = False,
106         inv_iterations: int = 6,  # recommended default in paper was 6.
107         v_skip_connection: Optional[nn.Module] = None,
108         conv_kernel_size: Optional[int] = None,
109         *args,
110         **kwargs,
111     ):
112         """
113         Nystrom attention mechanism, from Nystromformer_.
114         ::
115 
116             "A Nystrom-based Algorithm for Approximating Self-Attention."
117             Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., Singh, V. (2021)
118 
119             Reference codebase: https://github.com/mlpen/Nystromformer
120 
121         .. _Nystromformer: https://arxiv.org/pdf/2102.03902.pdf
122 
123         """
124         super().__init__()
125         # merged key padding mask and attention mask is not accepted
126         self.requires_separate_masks = True
127         self.num_landmarks = num_landmarks
128         # TODO: should be able to not have to pass in num_heads
129         self.num_heads = num_heads
130         self.use_razavi_pinverse = use_razavi_pinverse
131         self.pinverse_original_init = pinverse_original_init
132         self.inv_iterations = inv_iterations
133         self.attn_drop = nn.Dropout(dropout)
134         self.skip_connection = v_skip_connection
135         self.causal = causal
136 
137         if self.skip_connection is None and conv_kernel_size is not None:
138             self.skip_connection = nn.Conv2d(
139                 in_channels=self.num_heads,
140                 out_channels=self.num_heads,
141                 kernel_size=(conv_kernel_size, 1),
142                 padding=(conv_kernel_size // 2, 0),
143                 bias=False,
144                 groups=self.num_heads,
145             )
146 
147         if landmark_pooling is not None:
148             self.landmark_pooling = landmark_pooling
149         else:
150             self.landmark_pooling = get_avg_pool(self.num_landmarks)
151 
152         # Optional lower triangular masks for causal attention
153         self.causal_mask_1: Optional[torch.Tensor] = None
154         self.causal_mask_2: Optional[torch.Tensor] = None
155         self.causal_mask_3: Optional[torch.Tensor] = None
156 
157         # This attention does not support attention masks
158         self.supports_attention_mask = False
159         self.supports_key_padding_mask = True
160 
161     def forward(
162         self,
163         q: torch.Tensor,
164         k: torch.Tensor,
165         v: torch.Tensor,
166         key_padding_mask: Optional[torch.Tensor] = None,
167         *args,
168         **kwargs,
169     ):
170         r"""
171         key_padding_mask    Only a key padding mask is accepted here. The size must be (batch size, sequence length) or
172                             (batch size * num_heads, 1, sequence length). If dimensions are not correct, the mask will
173                             be ignored. An additive mask is expected, meaning float values using "-inf" to mask values
174         """
175 
176         batched_dim = k.size(0)
177         seq_len = k.size(-2)
178         tt = {"dtype": q.dtype, "device": q.device}
179 
180         if key_padding_mask is not None:
181             if key_padding_mask.dtype == torch.bool:
182                 logging.warning(
183                     "Bool mask found, but an additive mask is expected. Converting but this is slow"
184                 )
185                 key_padding_mask = bool_mask_to_additive(key_padding_mask)
186 
187             if key_padding_mask.ndim == 2:
188                 key_padding_mask = reshape_key_padding_mask(
189                     key_padding_mask, batched_dim
190                 )
191 
192             assert key_padding_mask.size() == (batched_dim, 1, seq_len), (
193                 f"key_padding_mask has invalid dimensions {key_padding_mask.size()}."
194                 f" Must have dimensions {batched_dim, 1, seq_len} or (batch_size, {seq_len})."
195             )
196 
197         if self.num_landmarks >= seq_len:
198             mask: Optional[torch.Tensor] = None
199 
200             if self.causal:
201                 mask = self._triu_mask(batched_dim, seq_len, seq_len, **tt)
202 
203             if key_padding_mask is not None:
204                 mask = key_padding_mask if mask is None else mask + key_padding_mask
205 
206             x = scaled_dot_product_attention(q=q, k=k, v=v, att_mask=mask)
207 
208         else:
209             q_landmarks = self.landmark_pooling(q)
210             k_landmarks = self.landmark_pooling(k)
211 
212             if self.causal and (
213                 self.causal_mask_1 is None
214                 or (batched_dim, seq_len, self.num_landmarks)
215                 != self.causal_mask_1.size()
216             ):
217                 self.causal_mask_1 = self._triu_mask(
218                     batched_dim, seq_len, self.num_landmarks, **tt
219                 )
220                 self.causal_mask_2 = self._triu_mask(
221                     batched_dim, self.num_landmarks, self.num_landmarks, **tt
222                 )
223                 self.causal_mask_3 = self._triu_mask(
224                     batched_dim, self.num_landmarks, seq_len, **tt
225                 )
226 
227             mask_1: Optional[torch.Tensor] = self.causal_mask_1
228             mask_2: Optional[torch.Tensor] = self.causal_mask_2
229             mask_3: Optional[torch.Tensor] = self.causal_mask_3
230             if key_padding_mask is not None:
231                 mask_1 = (
232                     key_padding_mask.transpose(-2, -1)
233                     if mask_1 is None
234                     else mask_1 + key_padding_mask.transpose(-2, -1)
235                 )
236                 mask_3 = (
237                     key_padding_mask if mask_3 is None else mask_3 + key_padding_mask
238                 )
239 
240             kernel_1 = scaled_query_key_softmax(q=q, k=k_landmarks, att_mask=mask_1)
241             kernel_2 = scaled_query_key_softmax(
242                 q=q_landmarks, k=k_landmarks, att_mask=mask_2
243             )
244             kernel_3 = scaled_dot_product_attention(
245                 q=q_landmarks, k=k, v=v, att_mask=mask_3
246             )
247 
248             kernel_2_inv = (
249                 iterative_pinv(
250                     kernel_2, self.inv_iterations, self.pinverse_original_init
251                 )
252                 if self.use_razavi_pinverse
253                 else torch.linalg.pinv(kernel_2)
254             )
255 
256             x = torch.matmul(
257                 torch.matmul(
258                     kernel_1,
259                     kernel_2_inv,
260                 ),
261                 kernel_3,
262             )
263 
264         if self.skip_connection:
265             # Assumption here is that v is 3D.
266             v_conv = self.skip_connection(
267                 v.reshape(-1, self.num_heads, v.size(-2), v.size(-1))
268             )
269             x += v_conv.reshape(-1, v_conv.size(-2), v_conv.size(-1))
270         x = self.attn_drop(x)
271         return x
272 
273     def _triu_mask(self, dim_1: int, dim_2: int, dim_3: int, **kwargs) -> torch.Tensor:
274         device = kwargs["device"]
275         dtype = kwargs["dtype"]
276 
277         return torch.triu(
278             torch.ones(dim_2, dim_3, dtype=dtype, device=device) * float("-inf"),
279             diagonal=1,
280         ).expand(
281             dim_1, -1, -1
282         )  # micro optim, save memory on the batch dimension
283
```
--- 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/xformers/components/attention/nystrom.py b/xformers/components/attention/nystrom.py
--- a/xformers/components/attention/nystrom.py
+++ b/xformers/components/attention/nystrom.py
@@ -57,24 +57,29 @@
     use_razavi_pinverse: Optional[bool]
 
 
-def get_avg_pool(n: int):
-    def avg_pool(x: torch.Tensor):
+class AvgPool(nn.Module):
+    def __init__(self, n: int):
+        super().__init__()
+        self.n = n
+
+    def forward(self, x: torch.Tensor):
         # Average independently for every segment in the sequence dimension
         seq_len = x.shape[1]
         head_dim = x.shape[2]
-        segments = seq_len // n
+        segments = seq_len // self.n
+        assert segments > 0, "num_landmarks should be smaller than the sequence length"
 
         # Dimensions are a match
-        if seq_len % n == 0:
+        if seq_len % self.n == 0:
             return x.reshape(
                 -1,
-                n,
+                self.n,
                 segments,
                 head_dim,
             ).mean(dim=-2)
 
         # Handle the last segment boundary being off
-        n_round = n - seq_len % n
+        n_round = self.n - seq_len % self.n
 
         x_avg_round = (
             x[:, : n_round * segments, :]
@@ -83,13 +88,11 @@
         )
         x_avg_off = (
             x[:, n_round * segments :, :]
-            .reshape(-1, n - n_round, segments + 1, head_dim)
+            .reshape(-1, self.n - n_round, segments + 1, head_dim)
             .mean(dim=-2)
         )
         return torch.cat((x_avg_round, x_avg_off), dim=-2)
 
-    return avg_pool
-
 
 @register_attention("nystrom", NystromSelfAttentionConfig)
 class NystromAttention(Attention):
@@ -147,7 +150,7 @@
         if landmark_pooling is not None:
             self.landmark_pooling = landmark_pooling
         else:
-            self.landmark_pooling = get_avg_pool(self.num_landmarks)
+            self.landmark_pooling = AvgPool(n=self.num_landmarks)
 
         # Optional lower triangular masks for causal attention
         self.causal_mask_1: Optional[torch.Tensor] = None

{"golden_diff": "diff --git a/xformers/components/attention/nystrom.py b/xformers/components/attention/nystrom.py\n--- a/xformers/components/attention/nystrom.py\n+++ b/xformers/components/attention/nystrom.py\n@@ -57,24 +57,29 @@\n     use_razavi_pinverse: Optional[bool]\n \n \n-def get_avg_pool(n: int):\n-    def avg_pool(x: torch.Tensor):\n+class AvgPool(nn.Module):\n+    def __init__(self, n: int):\n+        super().__init__()\n+        self.n = n\n+\n+    def forward(self, x: torch.Tensor):\n         # Average independently for every segment in the sequence dimension\n         seq_len = x.shape[1]\n         head_dim = x.shape[2]\n-        segments = seq_len // n\n+        segments = seq_len // self.n\n+        assert segments > 0, \"num_landmarks should be smaller than the sequence length\"\n \n         # Dimensions are a match\n-        if seq_len % n == 0:\n+        if seq_len % self.n == 0:\n             return x.reshape(\n                 -1,\n-                n,\n+                self.n,\n                 segments,\n                 head_dim,\n             ).mean(dim=-2)\n \n         # Handle the last segment boundary being off\n-        n_round = n - seq_len % n\n+        n_round = self.n - seq_len % self.n\n \n         x_avg_round = (\n             x[:, : n_round * segments, :]\n@@ -83,13 +88,11 @@\n         )\n         x_avg_off = (\n             x[:, n_round * segments :, :]\n-            .reshape(-1, n - n_round, segments + 1, head_dim)\n+            .reshape(-1, self.n - n_round, segments + 1, head_dim)\n             .mean(dim=-2)\n         )\n         return torch.cat((x_avg_round, x_avg_off), dim=-2)\n \n-    return avg_pool\n-\n \n @register_attention(\"nystrom\", NystromSelfAttentionConfig)\n class NystromAttention(Attention):\n@@ -147,7 +150,7 @@\n         if landmark_pooling is not None:\n             self.landmark_pooling = landmark_pooling\n         else:\n-            self.landmark_pooling = get_avg_pool(self.num_landmarks)\n+            self.landmark_pooling = AvgPool(n=self.num_landmarks)\n \n         # Optional lower triangular masks for causal attention\n         self.causal_mask_1: Optional[torch.Tensor] = None\n", "issue": "Cannot pickle nystrom based model\nWhen I `torch.save` my model with `nystrom` attention I get the following pickle error:\r\n\r\n```py\r\nAttributeError: Can't pickle local object 'get_avg_pool.<locals>.avg_pool'\r\n```\r\n\r\nI believe coming from this function:\r\n\r\nhttps://github.com/facebookresearch/xformers/blob/9232b2d27a775a43173e2fd86c03251ab64f7ede/xformers/components/attention/nystrom.py#L60\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport logging\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\n\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\nfrom xformers.components.attention.core import (\n    scaled_dot_product_attention,\n    scaled_query_key_softmax,\n)\nfrom xformers.components.attention.utils import (\n    bool_mask_to_additive,\n    iterative_pinv,\n    reshape_key_padding_mask,\n)\n\n\n@dataclass\nclass NystromSelfAttentionConfig(AttentionConfig):\n    \"\"\"\n    num_heads               Number of heads.\n    num_landmarks           Number of landmarks to use for softmax approximation. 64 often sufficient for a good\n                            approximation according to https://arxiv.org/pdf/2102.03902.pdf.\n    causal                  Apply a causal mask, in that the attention cannot be applied to the future.\n    use_razavi_pinverse     If true, use iterative method from (Razavi et al. 2014) to approximate the Moore-Penrose\n                            inverse, otherwise use standard torch inverse.\n    pinverse_original_init  True if using original initialization when calculating Moore-Penrose pseudo inverse using\n                            method from (Razavi et al. 2014).\n                            False if using exact coefficient computation (leads to faster convergence).\n    inv_iterations          Number of iterations for calculating the Moore-Penrose pseudo inverse.\n    v_skip_connection       A module that will take V as input and will be added as a skip connection to the\n                            softmax approximation. A skip connection is added in the paper to help with training.\n    conv_kernel_size        Kernel size for convolution optionally added to help in training.\n                            If v_skip_connection is not specified, this will be used to define the default\n                            depth wise convolution used as a skip connection.\n                            If both conv_kernel_size and v_skip_connection are None, no skip connection will\n                            be added.\n    landmark_pooling        Which module to use when computing landmarks. Default is AdaptiveAvgPool2d.\n    \"\"\"\n\n    num_heads: int\n    num_landmarks: Optional[int]\n    landmark_pooling: Optional[nn.Module]\n    causal: Optional[bool]\n    pinverse_original_init: Optional[bool]\n    inv_iterations: Optional[int]\n    v_skip_connection: Optional[nn.Module]\n    conv_kernel_size: Optional[int]\n    use_razavi_pinverse: Optional[bool]\n\n\ndef get_avg_pool(n: int):\n    def avg_pool(x: torch.Tensor):\n        # Average independently for every segment in the sequence dimension\n        seq_len = x.shape[1]\n        head_dim = x.shape[2]\n        segments = seq_len // n\n\n        # Dimensions are a match\n        if seq_len % n == 0:\n            return x.reshape(\n                -1,\n                n,\n                segments,\n                head_dim,\n            ).mean(dim=-2)\n\n        # Handle the last segment boundary being off\n        n_round = n - seq_len % n\n\n        x_avg_round = (\n            x[:, : n_round * segments, :]\n            .reshape(-1, n_round, segments, head_dim)\n            .mean(dim=-2)\n        )\n        x_avg_off = (\n            x[:, n_round * segments :, :]\n            .reshape(-1, n - n_round, segments + 1, head_dim)\n            .mean(dim=-2)\n        )\n        return torch.cat((x_avg_round, x_avg_off), dim=-2)\n\n    return avg_pool\n\n\n@register_attention(\"nystrom\", NystromSelfAttentionConfig)\nclass NystromAttention(Attention):\n    # TODO: update defaults for use_razavi_pinverse and inv_iterations\n    def __init__(\n        self,\n        dropout: float,\n        num_heads: int,\n        num_landmarks: int = 64,\n        landmark_pooling: Optional[nn.Module] = None,\n        causal: bool = False,\n        use_razavi_pinverse: bool = True,\n        pinverse_original_init: bool = False,\n        inv_iterations: int = 6,  # recommended default in paper was 6.\n        v_skip_connection: Optional[nn.Module] = None,\n        conv_kernel_size: Optional[int] = None,\n        *args,\n        **kwargs,\n    ):\n        \"\"\"\n        Nystrom attention mechanism, from Nystromformer_.\n        ::\n\n            \"A Nystrom-based Algorithm for Approximating Self-Attention.\"\n            Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., Singh, V. (2021)\n\n            Reference codebase: https://github.com/mlpen/Nystromformer\n\n        .. _Nystromformer: https://arxiv.org/pdf/2102.03902.pdf\n\n        \"\"\"\n        super().__init__()\n        # merged key padding mask and attention mask is not accepted\n        self.requires_separate_masks = True\n        self.num_landmarks = num_landmarks\n        # TODO: should be able to not have to pass in num_heads\n        self.num_heads = num_heads\n        self.use_razavi_pinverse = use_razavi_pinverse\n        self.pinverse_original_init = pinverse_original_init\n        self.inv_iterations = inv_iterations\n        self.attn_drop = nn.Dropout(dropout)\n        self.skip_connection = v_skip_connection\n        self.causal = causal\n\n        if self.skip_connection is None and conv_kernel_size is not None:\n            self.skip_connection = nn.Conv2d(\n                in_channels=self.num_heads,\n                out_channels=self.num_heads,\n                kernel_size=(conv_kernel_size, 1),\n                padding=(conv_kernel_size // 2, 0),\n                bias=False,\n                groups=self.num_heads,\n            )\n\n        if landmark_pooling is not None:\n            self.landmark_pooling = landmark_pooling\n        else:\n            self.landmark_pooling = get_avg_pool(self.num_landmarks)\n\n        # Optional lower triangular masks for causal attention\n        self.causal_mask_1: Optional[torch.Tensor] = None\n        self.causal_mask_2: Optional[torch.Tensor] = None\n        self.causal_mask_3: Optional[torch.Tensor] = None\n\n        # This attention does not support attention masks\n        self.supports_attention_mask = False\n        self.supports_key_padding_mask = True\n\n    def forward(\n        self,\n        q: torch.Tensor,\n        k: torch.Tensor,\n        v: torch.Tensor,\n        key_padding_mask: Optional[torch.Tensor] = None,\n        *args,\n        **kwargs,\n    ):\n        r\"\"\"\n        key_padding_mask    Only a key padding mask is accepted here. The size must be (batch size, sequence length) or\n                            (batch size * num_heads, 1, sequence length). If dimensions are not correct, the mask will\n                            be ignored. An additive mask is expected, meaning float values using \"-inf\" to mask values\n        \"\"\"\n\n        batched_dim = k.size(0)\n        seq_len = k.size(-2)\n        tt = {\"dtype\": q.dtype, \"device\": q.device}\n\n        if key_padding_mask is not None:\n            if key_padding_mask.dtype == torch.bool:\n                logging.warning(\n                    \"Bool mask found, but an additive mask is expected. Converting but this is slow\"\n                )\n                key_padding_mask = bool_mask_to_additive(key_padding_mask)\n\n            if key_padding_mask.ndim == 2:\n                key_padding_mask = reshape_key_padding_mask(\n                    key_padding_mask, batched_dim\n                )\n\n            assert key_padding_mask.size() == (batched_dim, 1, seq_len), (\n                f\"key_padding_mask has invalid dimensions {key_padding_mask.size()}.\"\n                f\" Must have dimensions {batched_dim, 1, seq_len} or (batch_size, {seq_len}).\"\n            )\n\n        if self.num_landmarks >= seq_len:\n            mask: Optional[torch.Tensor] = None\n\n            if self.causal:\n                mask = self._triu_mask(batched_dim, seq_len, seq_len, **tt)\n\n            if key_padding_mask is not None:\n                mask = key_padding_mask if mask is None else mask + key_padding_mask\n\n            x = scaled_dot_product_attention(q=q, k=k, v=v, att_mask=mask)\n\n        else:\n            q_landmarks = self.landmark_pooling(q)\n            k_landmarks = self.landmark_pooling(k)\n\n            if self.causal and (\n                self.causal_mask_1 is None\n                or (batched_dim, seq_len, self.num_landmarks)\n                != self.causal_mask_1.size()\n            ):\n                self.causal_mask_1 = self._triu_mask(\n                    batched_dim, seq_len, self.num_landmarks, **tt\n                )\n                self.causal_mask_2 = self._triu_mask(\n                    batched_dim, self.num_landmarks, self.num_landmarks, **tt\n                )\n                self.causal_mask_3 = self._triu_mask(\n                    batched_dim, self.num_landmarks, seq_len, **tt\n                )\n\n            mask_1: Optional[torch.Tensor] = self.causal_mask_1\n            mask_2: Optional[torch.Tensor] = self.causal_mask_2\n            mask_3: Optional[torch.Tensor] = self.causal_mask_3\n            if key_padding_mask is not None:\n                mask_1 = (\n                    key_padding_mask.transpose(-2, -1)\n                    if mask_1 is None\n                    else mask_1 + key_padding_mask.transpose(-2, -1)\n                )\n                mask_3 = (\n                    key_padding_mask if mask_3 is None else mask_3 + key_padding_mask\n                )\n\n            kernel_1 = scaled_query_key_softmax(q=q, k=k_landmarks, att_mask=mask_1)\n            kernel_2 = scaled_query_key_softmax(\n                q=q_landmarks, k=k_landmarks, att_mask=mask_2\n            )\n            kernel_3 = scaled_dot_product_attention(\n                q=q_landmarks, k=k, v=v, att_mask=mask_3\n            )\n\n            kernel_2_inv = (\n                iterative_pinv(\n                    kernel_2, self.inv_iterations, self.pinverse_original_init\n                )\n                if self.use_razavi_pinverse\n                else torch.linalg.pinv(kernel_2)\n            )\n\n            x = torch.matmul(\n                torch.matmul(\n                    kernel_1,\n                    kernel_2_inv,\n                ),\n                kernel_3,\n            )\n\n        if self.skip_connection:\n            # Assumption here is that v is 3D.\n            v_conv = self.skip_connection(\n                v.reshape(-1, self.num_heads, v.size(-2), v.size(-1))\n            )\n            x += v_conv.reshape(-1, v_conv.size(-2), v_conv.size(-1))\n        x = self.attn_drop(x)\n        return x\n\n    def _triu_mask(self, dim_1: int, dim_2: int, dim_3: int, **kwargs) -> torch.Tensor:\n        device = kwargs[\"device\"]\n        dtype = kwargs[\"dtype\"]\n\n        return torch.triu(\n            torch.ones(dim_2, dim_3, dtype=dtype, device=device) * float(\"-inf\"),\n            diagonal=1,\n        ).expand(\n            dim_1, -1, -1\n        )  # micro optim, save memory on the batch dimension\n", "path": "xformers/components/attention/nystrom.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport logging\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\n\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\nfrom xformers.components.attention.core import (\n    scaled_dot_product_attention,\n    scaled_query_key_softmax,\n)\nfrom xformers.components.attention.utils import (\n    bool_mask_to_additive,\n    iterative_pinv,\n    reshape_key_padding_mask,\n)\n\n\n@dataclass\nclass NystromSelfAttentionConfig(AttentionConfig):\n    \"\"\"\n    num_heads               Number of heads.\n    num_landmarks           Number of landmarks to use for softmax approximation. 64 often sufficient for a good\n                            approximation according to https://arxiv.org/pdf/2102.03902.pdf.\n    causal                  Apply a causal mask, in that the attention cannot be applied to the future.\n    use_razavi_pinverse     If true, use iterative method from (Razavi et al. 2014) to approximate the Moore-Penrose\n                            inverse, otherwise use standard torch inverse.\n    pinverse_original_init  True if using original initialization when calculating Moore-Penrose pseudo inverse using\n                            method from (Razavi et al. 2014).\n                            False if using exact coefficient computation (leads to faster convergence).\n    inv_iterations          Number of iterations for calculating the Moore-Penrose pseudo inverse.\n    v_skip_connection       A module that will take V as input and will be added as a skip connection to the\n                            softmax approximation. A skip connection is added in the paper to help with training.\n    conv_kernel_size        Kernel size for convolution optionally added to help in training.\n                            If v_skip_connection is not specified, this will be used to define the default\n                            depth wise convolution used as a skip connection.\n                            If both conv_kernel_size and v_skip_connection are None, no skip connection will\n                            be added.\n    landmark_pooling        Which module to use when computing landmarks. Default is AdaptiveAvgPool2d.\n    \"\"\"\n\n    num_heads: int\n    num_landmarks: Optional[int]\n    landmark_pooling: Optional[nn.Module]\n    causal: Optional[bool]\n    pinverse_original_init: Optional[bool]\n    inv_iterations: Optional[int]\n    v_skip_connection: Optional[nn.Module]\n    conv_kernel_size: Optional[int]\n    use_razavi_pinverse: Optional[bool]\n\n\nclass AvgPool(nn.Module):\n    def __init__(self, n: int):\n        super().__init__()\n        self.n = n\n\n    def forward(self, x: torch.Tensor):\n        # Average independently for every segment in the sequence dimension\n        seq_len = x.shape[1]\n        head_dim = x.shape[2]\n        segments = seq_len // self.n\n        assert segments > 0, \"num_landmarks should be smaller than the sequence length\"\n\n        # Dimensions are a match\n        if seq_len % self.n == 0:\n            return x.reshape(\n                -1,\n                self.n,\n                segments,\n                head_dim,\n            ).mean(dim=-2)\n\n        # Handle the last segment boundary being off\n        n_round = self.n - seq_len % self.n\n\n        x_avg_round = (\n            x[:, : n_round * segments, :]\n            .reshape(-1, n_round, segments, head_dim)\n            .mean(dim=-2)\n        )\n        x_avg_off = (\n            x[:, n_round * segments :, :]\n            .reshape(-1, self.n - n_round, segments + 1, head_dim)\n            .mean(dim=-2)\n        )\n        return torch.cat((x_avg_round, x_avg_off), dim=-2)\n\n\n@register_attention(\"nystrom\", NystromSelfAttentionConfig)\nclass NystromAttention(Attention):\n    # TODO: update defaults for use_razavi_pinverse and inv_iterations\n    def __init__(\n        self,\n        dropout: float,\n        num_heads: int,\n        num_landmarks: int = 64,\n        landmark_pooling: Optional[nn.Module] = None,\n        causal: bool = False,\n        use_razavi_pinverse: bool = True,\n        pinverse_original_init: bool = False,\n        inv_iterations: int = 6,  # recommended default in paper was 6.\n        v_skip_connection: Optional[nn.Module] = None,\n        conv_kernel_size: Optional[int] = None,\n        *args,\n        **kwargs,\n    ):\n        \"\"\"\n        Nystrom attention mechanism, from Nystromformer_.\n        ::\n\n            \"A Nystrom-based Algorithm for Approximating Self-Attention.\"\n            Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., Singh, V. (2021)\n\n            Reference codebase: https://github.com/mlpen/Nystromformer\n\n        .. _Nystromformer: https://arxiv.org/pdf/2102.03902.pdf\n\n        \"\"\"\n        super().__init__()\n        # merged key padding mask and attention mask is not accepted\n        self.requires_separate_masks = True\n        self.num_landmarks = num_landmarks\n        # TODO: should be able to not have to pass in num_heads\n        self.num_heads = num_heads\n        self.use_razavi_pinverse = use_razavi_pinverse\n        self.pinverse_original_init = pinverse_original_init\n        self.inv_iterations = inv_iterations\n        self.attn_drop = nn.Dropout(dropout)\n        self.skip_connection = v_skip_connection\n        self.causal = causal\n\n        if self.skip_connection is None and conv_kernel_size is not None:\n            self.skip_connection = nn.Conv2d(\n                in_channels=self.num_heads,\n                out_channels=self.num_heads,\n                kernel_size=(conv_kernel_size, 1),\n                padding=(conv_kernel_size // 2, 0),\n                bias=False,\n                groups=self.num_heads,\n            )\n\n        if landmark_pooling is not None:\n            self.landmark_pooling = landmark_pooling\n        else:\n            self.landmark_pooling = AvgPool(n=self.num_landmarks)\n\n        # Optional lower triangular masks for causal attention\n        self.causal_mask_1: Optional[torch.Tensor] = None\n        self.causal_mask_2: Optional[torch.Tensor] = None\n        self.causal_mask_3: Optional[torch.Tensor] = None\n\n        # This attention does not support attention masks\n        self.supports_attention_mask = False\n        self.supports_key_padding_mask = True\n\n    def forward(\n        self,\n        q: torch.Tensor,\n        k: torch.Tensor,\n        v: torch.Tensor,\n        key_padding_mask: Optional[torch.Tensor] = None,\n        *args,\n        **kwargs,\n    ):\n        r\"\"\"\n        key_padding_mask    Only a key padding mask is accepted here. The size must be (batch size, sequence length) or\n                            (batch size * num_heads, 1, sequence length). If dimensions are not correct, the mask will\n                            be ignored. An additive mask is expected, meaning float values using \"-inf\" to mask values\n        \"\"\"\n\n        batched_dim = k.size(0)\n        seq_len = k.size(-2)\n        tt = {\"dtype\": q.dtype, \"device\": q.device}\n\n        if key_padding_mask is not None:\n            if key_padding_mask.dtype == torch.bool:\n                logging.warning(\n                    \"Bool mask found, but an additive mask is expected. Converting but this is slow\"\n                )\n                key_padding_mask = bool_mask_to_additive(key_padding_mask)\n\n            if key_padding_mask.ndim == 2:\n                key_padding_mask = reshape_key_padding_mask(\n                    key_padding_mask, batched_dim\n                )\n\n            assert key_padding_mask.size() == (batched_dim, 1, seq_len), (\n                f\"key_padding_mask has invalid dimensions {key_padding_mask.size()}.\"\n                f\" Must have dimensions {batched_dim, 1, seq_len} or (batch_size, {seq_len}).\"\n            )\n\n        if self.num_landmarks >= seq_len:\n            mask: Optional[torch.Tensor] = None\n\n            if self.causal:\n                mask = self._triu_mask(batched_dim, seq_len, seq_len, **tt)\n\n            if key_padding_mask is not None:\n                mask = key_padding_mask if mask is None else mask + key_padding_mask\n\n            x = scaled_dot_product_attention(q=q, k=k, v=v, att_mask=mask)\n\n        else:\n            q_landmarks = self.landmark_pooling(q)\n            k_landmarks = self.landmark_pooling(k)\n\n            if self.causal and (\n                self.causal_mask_1 is None\n                or (batched_dim, seq_len, self.num_landmarks)\n                != self.causal_mask_1.size()\n            ):\n                self.causal_mask_1 = self._triu_mask(\n                    batched_dim, seq_len, self.num_landmarks, **tt\n                )\n                self.causal_mask_2 = self._triu_mask(\n                    batched_dim, self.num_landmarks, self.num_landmarks, **tt\n                )\n                self.causal_mask_3 = self._triu_mask(\n                    batched_dim, self.num_landmarks, seq_len, **tt\n                )\n\n            mask_1: Optional[torch.Tensor] = self.causal_mask_1\n            mask_2: Optional[torch.Tensor] = self.causal_mask_2\n            mask_3: Optional[torch.Tensor] = self.causal_mask_3\n            if key_padding_mask is not None:\n                mask_1 = (\n                    key_padding_mask.transpose(-2, -1)\n                    if mask_1 is None\n                    else mask_1 + key_padding_mask.transpose(-2, -1)\n                )\n                mask_3 = (\n                    key_padding_mask if mask_3 is None else mask_3 + key_padding_mask\n                )\n\n            kernel_1 = scaled_query_key_softmax(q=q, k=k_landmarks, att_mask=mask_1)\n            kernel_2 = scaled_query_key_softmax(\n                q=q_landmarks, k=k_landmarks, att_mask=mask_2\n            )\n            kernel_3 = scaled_dot_product_attention(\n                q=q_landmarks, k=k, v=v, att_mask=mask_3\n            )\n\n            kernel_2_inv = (\n                iterative_pinv(\n                    kernel_2, self.inv_iterations, self.pinverse_original_init\n                )\n                if self.use_razavi_pinverse\n                else torch.linalg.pinv(kernel_2)\n            )\n\n            x = torch.matmul(\n                torch.matmul(\n                    kernel_1,\n                    kernel_2_inv,\n                ),\n                kernel_3,\n            )\n\n        if self.skip_connection:\n            # Assumption here is that v is 3D.\n            v_conv = self.skip_connection(\n                v.reshape(-1, self.num_heads, v.size(-2), v.size(-1))\n            )\n            x += v_conv.reshape(-1, v_conv.size(-2), v_conv.size(-1))\n        x = self.attn_drop(x)\n        return x\n\n    def _triu_mask(self, dim_1: int, dim_2: int, dim_3: int, **kwargs) -> torch.Tensor:\n        device = kwargs[\"device\"]\n        dtype = kwargs[\"dtype\"]\n\n        return torch.triu(\n            torch.ones(dim_2, dim_3, dtype=dtype, device=device) * float(\"-inf\"),\n            diagonal=1,\n        ).expand(\n            dim_1, -1, -1\n        )  # micro optim, save memory on the batch dimension\n", "path": "xformers/components/attention/nystrom.py"}]}