We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
JIT scripting is broken
# 🐛 Bug
JIT scripting xformers (running commit 357545ae13948659db07428553155e1802ee15af)  breaks with the following error:

```bash
xformers/components/attention/attention_mask.py", line 128
    def __add__(self, other):
        assert isinstance(other, type(self))
                                 ~~~~~~~~~ <--- HERE
        return AttentionMask(self.values + other.values, is_causal=False)
'AttentionMask.__add__' is being compiled since it was called from '__torch__.xformers.components.attention.attention_mask.AttentionMask'
```

## To Reproduce

Blocks config:
```yaml
reversible: false
block_type: "encoder"
num_layers: 12
dim_model: 768
layer_norm_style: "pre"

multi_head_config:
  num_heads: 12
  residual_dropout: 0.0
  use_rotary_embeddings: false

  attention:
    name: "scaled_dot_product"
    dropout: 0.0
    causal: false

feedforward_config:
  name: "MLP"
  dropout: 0.0
  activation: "gelu"
  hidden_layer_multiplier: 4
```

Python code to repro:
```python
import yaml
import torch

from xformers.factory import xFormer, xFormerConfig

with open(xformer_config_file, "rb") as fileptr:  # above config
    model_config = yaml.load(fileptr, Loader=yaml.FullLoader)

torch.jit.script(xFormer.from_config(xFormerConfig([model_config])))
```


--- 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/attention_mask.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 from typing import Optional, Type, TypeVar
8 
9 import torch
10 
11 Self = TypeVar("Self", bound="AttentionMask")
12 
13 
14 class AttentionMask:
15     """
16     Holds an attention mask, along with a couple of helpers and attributes.
17 
18     .. note: this is an additive mask, meaning that coefficients which should be computed hold the '0.' value,
19         and coefficients which should be skipped hold the '-inf' value. Any other value is possible if the purpose
20         is to bias the attention computation for instance
21 
22     .. note: the attention mask dimensions are expected to be `[batch, to_sequence, from_sequence]`,
23         `[to_sequence, from_sequence]`, or anything broadcastable in between
24     """
25 
26     def __init__(self, additive_mask: torch.Tensor, is_causal: bool = False):
27         assert additive_mask.is_floating_point(), additive_mask.dtype
28         assert not additive_mask.requires_grad
29 
30         if additive_mask.ndim == 2:
31             additive_mask = additive_mask.unsqueeze(0)
32 
33         self.values = additive_mask
34         self.is_causal = is_causal
35         self.seq_len = additive_mask.shape[1]
36         self.to_seq_len = additive_mask.shape[0]
37 
38     def to_bool(self) -> torch.Tensor:
39         """
40         .. warning: we assume here that True implies that the value should be computed
41         """
42         return self.values != float("-inf")
43 
44     @classmethod
45     def from_bool(cls: Type[Self], x: torch.Tensor) -> Self:
46         """
47         Create an AttentionMask given a boolean pattern.
48         .. warning: we assume here that True implies that the value should be computed
49         """
50         assert x.dtype == torch.bool
51 
52         additive_mask = torch.empty_like(x, dtype=torch.float, device=x.device)
53         additive_mask.masked_fill_(x, 0.0)
54         additive_mask.masked_fill_(~x, float("-inf"))
55 
56         return cls(additive_mask)
57 
58     @classmethod
59     def from_multiplicative(cls: Type[Self], x: torch.Tensor) -> Self:
60         """
61         Create an AttentionMask given a multiplicative attention mask.
62         """
63         assert not x.dtype == torch.bool
64 
65         additive_mask = torch.empty_like(x, dtype=torch.float, device=x.device)
66         x = x.bool()
67 
68         additive_mask.masked_fill_(x, 0.0)
69         additive_mask.masked_fill_(~x, float("-inf"))
70 
71         return cls(additive_mask)
72 
73     @classmethod
74     def make_causal(
75         cls: Type[Self],
76         seq_len: int,
77         to_seq_len: Optional[int] = None,
78         device: Optional[torch.device] = None,
79         dtype: Optional[torch.dtype] = None,
80     ) -> Self:
81         if not to_seq_len:
82             to_seq_len = seq_len
83 
84         additive_mask = torch.triu(
85             torch.ones(seq_len, to_seq_len, device=device, dtype=dtype) * float("-inf"),
86             diagonal=1,
87         )
88         return cls(additive_mask=additive_mask, is_causal=True)
89 
90     def make_crop(
91         self, seq_len: int, to_seq_len: Optional[int] = None
92     ) -> "AttentionMask":
93         """
94         Return a cropped attention mask, whose underlying tensor is a view of this one
95         """
96 
97         if not to_seq_len:
98             to_seq_len = seq_len
99 
100         return AttentionMask(
101             self.values[:, :seq_len, :to_seq_len], is_causal=self.is_causal
102         )
103 
104     def __repr__(self):
105         return f"AttentionMask - causal {self.is_causal} - mask " + str(self.values)
106 
107     @property
108     def device(self):
109         return self.values.device
110 
111     @property
112     def is_sparse(self):
113         return False
114 
115     @property
116     def ndim(self):
117         return len(self.values.shape)
118 
119     @property
120     def dtype(self):
121         return self.values.dtype
122 
123     @property
124     def shape(self):
125         return self.values.shape
126 
127     def __add__(self, other):
128         assert isinstance(other, type(self))
129         return AttentionMask(self.values + other.values, is_causal=False)
130 
131     def to(
132         self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None
133     ) -> "AttentionMask":
134         assert device is None or isinstance(device, torch.device)
135         assert dtype is None or isinstance(dtype, torch.dtype)
136         assert device is not None or dtype is not None
137 
138         # Noop if we don't need to create another instance
139         if ((device and device == self.device) or not device) and (
140             (dtype and dtype == self.dtype) or not dtype
141         ):
142             return self
143 
144         return AttentionMask(self.values.to(device=device, dtype=dtype), self.is_causal)
145
```
--- 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/attention_mask.py b/xformers/components/attention/attention_mask.py
--- a/xformers/components/attention/attention_mask.py
+++ b/xformers/components/attention/attention_mask.py
@@ -125,7 +125,6 @@
         return self.values.shape
 
     def __add__(self, other):
-        assert isinstance(other, type(self))
         return AttentionMask(self.values + other.values, is_causal=False)
 
     def to(

{"golden_diff": "diff --git a/xformers/components/attention/attention_mask.py b/xformers/components/attention/attention_mask.py\n--- a/xformers/components/attention/attention_mask.py\n+++ b/xformers/components/attention/attention_mask.py\n@@ -125,7 +125,6 @@\n         return self.values.shape\n \n     def __add__(self, other):\n-        assert isinstance(other, type(self))\n         return AttentionMask(self.values + other.values, is_causal=False)\n \n     def to(\n", "issue": "JIT scripting is broken\n# \ud83d\udc1b Bug\r\nJIT scripting xformers (running commit 357545ae13948659db07428553155e1802ee15af)  breaks with the following error:\r\n\r\n```bash\r\nxformers/components/attention/attention_mask.py\", line 128\r\n    def __add__(self, other):\r\n        assert isinstance(other, type(self))\r\n                                 ~~~~~~~~~ <--- HERE\r\n        return AttentionMask(self.values + other.values, is_causal=False)\r\n'AttentionMask.__add__' is being compiled since it was called from '__torch__.xformers.components.attention.attention_mask.AttentionMask'\r\n```\r\n\r\n## To Reproduce\r\n\r\nBlocks config:\r\n```yaml\r\nreversible: false\r\nblock_type: \"encoder\"\r\nnum_layers: 12\r\ndim_model: 768\r\nlayer_norm_style: \"pre\"\r\n\r\nmulti_head_config:\r\n  num_heads: 12\r\n  residual_dropout: 0.0\r\n  use_rotary_embeddings: false\r\n\r\n  attention:\r\n    name: \"scaled_dot_product\"\r\n    dropout: 0.0\r\n    causal: false\r\n\r\nfeedforward_config:\r\n  name: \"MLP\"\r\n  dropout: 0.0\r\n  activation: \"gelu\"\r\n  hidden_layer_multiplier: 4\r\n```\r\n\r\nPython code to repro:\r\n```python\r\nimport yaml\r\nimport torch\r\n\r\nfrom xformers.factory import xFormer, xFormerConfig\r\n\r\nwith open(xformer_config_file, \"rb\") as fileptr:  # above config\r\n    model_config = yaml.load(fileptr, Loader=yaml.FullLoader)\r\n\r\ntorch.jit.script(xFormer.from_config(xFormerConfig([model_config])))\r\n```\r\n\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\nfrom typing import Optional, Type, TypeVar\n\nimport torch\n\nSelf = TypeVar(\"Self\", bound=\"AttentionMask\")\n\n\nclass AttentionMask:\n    \"\"\"\n    Holds an attention mask, along with a couple of helpers and attributes.\n\n    .. note: this is an additive mask, meaning that coefficients which should be computed hold the '0.' value,\n        and coefficients which should be skipped hold the '-inf' value. Any other value is possible if the purpose\n        is to bias the attention computation for instance\n\n    .. note: the attention mask dimensions are expected to be `[batch, to_sequence, from_sequence]`,\n        `[to_sequence, from_sequence]`, or anything broadcastable in between\n    \"\"\"\n\n    def __init__(self, additive_mask: torch.Tensor, is_causal: bool = False):\n        assert additive_mask.is_floating_point(), additive_mask.dtype\n        assert not additive_mask.requires_grad\n\n        if additive_mask.ndim == 2:\n            additive_mask = additive_mask.unsqueeze(0)\n\n        self.values = additive_mask\n        self.is_causal = is_causal\n        self.seq_len = additive_mask.shape[1]\n        self.to_seq_len = additive_mask.shape[0]\n\n    def to_bool(self) -> torch.Tensor:\n        \"\"\"\n        .. warning: we assume here that True implies that the value should be computed\n        \"\"\"\n        return self.values != float(\"-inf\")\n\n    @classmethod\n    def from_bool(cls: Type[Self], x: torch.Tensor) -> Self:\n        \"\"\"\n        Create an AttentionMask given a boolean pattern.\n        .. warning: we assume here that True implies that the value should be computed\n        \"\"\"\n        assert x.dtype == torch.bool\n\n        additive_mask = torch.empty_like(x, dtype=torch.float, device=x.device)\n        additive_mask.masked_fill_(x, 0.0)\n        additive_mask.masked_fill_(~x, float(\"-inf\"))\n\n        return cls(additive_mask)\n\n    @classmethod\n    def from_multiplicative(cls: Type[Self], x: torch.Tensor) -> Self:\n        \"\"\"\n        Create an AttentionMask given a multiplicative attention mask.\n        \"\"\"\n        assert not x.dtype == torch.bool\n\n        additive_mask = torch.empty_like(x, dtype=torch.float, device=x.device)\n        x = x.bool()\n\n        additive_mask.masked_fill_(x, 0.0)\n        additive_mask.masked_fill_(~x, float(\"-inf\"))\n\n        return cls(additive_mask)\n\n    @classmethod\n    def make_causal(\n        cls: Type[Self],\n        seq_len: int,\n        to_seq_len: Optional[int] = None,\n        device: Optional[torch.device] = None,\n        dtype: Optional[torch.dtype] = None,\n    ) -> Self:\n        if not to_seq_len:\n            to_seq_len = seq_len\n\n        additive_mask = torch.triu(\n            torch.ones(seq_len, to_seq_len, device=device, dtype=dtype) * float(\"-inf\"),\n            diagonal=1,\n        )\n        return cls(additive_mask=additive_mask, is_causal=True)\n\n    def make_crop(\n        self, seq_len: int, to_seq_len: Optional[int] = None\n    ) -> \"AttentionMask\":\n        \"\"\"\n        Return a cropped attention mask, whose underlying tensor is a view of this one\n        \"\"\"\n\n        if not to_seq_len:\n            to_seq_len = seq_len\n\n        return AttentionMask(\n            self.values[:, :seq_len, :to_seq_len], is_causal=self.is_causal\n        )\n\n    def __repr__(self):\n        return f\"AttentionMask - causal {self.is_causal} - mask \" + str(self.values)\n\n    @property\n    def device(self):\n        return self.values.device\n\n    @property\n    def is_sparse(self):\n        return False\n\n    @property\n    def ndim(self):\n        return len(self.values.shape)\n\n    @property\n    def dtype(self):\n        return self.values.dtype\n\n    @property\n    def shape(self):\n        return self.values.shape\n\n    def __add__(self, other):\n        assert isinstance(other, type(self))\n        return AttentionMask(self.values + other.values, is_causal=False)\n\n    def to(\n        self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None\n    ) -> \"AttentionMask\":\n        assert device is None or isinstance(device, torch.device)\n        assert dtype is None or isinstance(dtype, torch.dtype)\n        assert device is not None or dtype is not None\n\n        # Noop if we don't need to create another instance\n        if ((device and device == self.device) or not device) and (\n            (dtype and dtype == self.dtype) or not dtype\n        ):\n            return self\n\n        return AttentionMask(self.values.to(device=device, dtype=dtype), self.is_causal)\n", "path": "xformers/components/attention/attention_mask.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\nfrom typing import Optional, Type, TypeVar\n\nimport torch\n\nSelf = TypeVar(\"Self\", bound=\"AttentionMask\")\n\n\nclass AttentionMask:\n    \"\"\"\n    Holds an attention mask, along with a couple of helpers and attributes.\n\n    .. note: this is an additive mask, meaning that coefficients which should be computed hold the '0.' value,\n        and coefficients which should be skipped hold the '-inf' value. Any other value is possible if the purpose\n        is to bias the attention computation for instance\n\n    .. note: the attention mask dimensions are expected to be `[batch, to_sequence, from_sequence]`,\n        `[to_sequence, from_sequence]`, or anything broadcastable in between\n    \"\"\"\n\n    def __init__(self, additive_mask: torch.Tensor, is_causal: bool = False):\n        assert additive_mask.is_floating_point(), additive_mask.dtype\n        assert not additive_mask.requires_grad\n\n        if additive_mask.ndim == 2:\n            additive_mask = additive_mask.unsqueeze(0)\n\n        self.values = additive_mask\n        self.is_causal = is_causal\n        self.seq_len = additive_mask.shape[1]\n        self.to_seq_len = additive_mask.shape[0]\n\n    def to_bool(self) -> torch.Tensor:\n        \"\"\"\n        .. warning: we assume here that True implies that the value should be computed\n        \"\"\"\n        return self.values != float(\"-inf\")\n\n    @classmethod\n    def from_bool(cls: Type[Self], x: torch.Tensor) -> Self:\n        \"\"\"\n        Create an AttentionMask given a boolean pattern.\n        .. warning: we assume here that True implies that the value should be computed\n        \"\"\"\n        assert x.dtype == torch.bool\n\n        additive_mask = torch.empty_like(x, dtype=torch.float, device=x.device)\n        additive_mask.masked_fill_(x, 0.0)\n        additive_mask.masked_fill_(~x, float(\"-inf\"))\n\n        return cls(additive_mask)\n\n    @classmethod\n    def from_multiplicative(cls: Type[Self], x: torch.Tensor) -> Self:\n        \"\"\"\n        Create an AttentionMask given a multiplicative attention mask.\n        \"\"\"\n        assert not x.dtype == torch.bool\n\n        additive_mask = torch.empty_like(x, dtype=torch.float, device=x.device)\n        x = x.bool()\n\n        additive_mask.masked_fill_(x, 0.0)\n        additive_mask.masked_fill_(~x, float(\"-inf\"))\n\n        return cls(additive_mask)\n\n    @classmethod\n    def make_causal(\n        cls: Type[Self],\n        seq_len: int,\n        to_seq_len: Optional[int] = None,\n        device: Optional[torch.device] = None,\n        dtype: Optional[torch.dtype] = None,\n    ) -> Self:\n        if not to_seq_len:\n            to_seq_len = seq_len\n\n        additive_mask = torch.triu(\n            torch.ones(seq_len, to_seq_len, device=device, dtype=dtype) * float(\"-inf\"),\n            diagonal=1,\n        )\n        return cls(additive_mask=additive_mask, is_causal=True)\n\n    def make_crop(\n        self, seq_len: int, to_seq_len: Optional[int] = None\n    ) -> \"AttentionMask\":\n        \"\"\"\n        Return a cropped attention mask, whose underlying tensor is a view of this one\n        \"\"\"\n\n        if not to_seq_len:\n            to_seq_len = seq_len\n\n        return AttentionMask(\n            self.values[:, :seq_len, :to_seq_len], is_causal=self.is_causal\n        )\n\n    def __repr__(self):\n        return f\"AttentionMask - causal {self.is_causal} - mask \" + str(self.values)\n\n    @property\n    def device(self):\n        return self.values.device\n\n    @property\n    def is_sparse(self):\n        return False\n\n    @property\n    def ndim(self):\n        return len(self.values.shape)\n\n    @property\n    def dtype(self):\n        return self.values.dtype\n\n    @property\n    def shape(self):\n        return self.values.shape\n\n    def __add__(self, other):\n        return AttentionMask(self.values + other.values, is_causal=False)\n\n    def to(\n        self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None\n    ) -> \"AttentionMask\":\n        assert device is None or isinstance(device, torch.device)\n        assert dtype is None or isinstance(dtype, torch.dtype)\n        assert device is not None or dtype is not None\n\n        # Noop if we don't need to create another instance\n        if ((device and device == self.device) or not device) and (\n            (dtype and dtype == self.dtype) or not dtype\n        ):\n            return self\n\n        return AttentionMask(self.values.to(device=device, dtype=dtype), self.is_causal)\n", "path": "xformers/components/attention/attention_mask.py"}]}