prompt: 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()
  
  ```
  
golden_diff: 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(
  
verification_info: {"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"}]}