We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[feat] Make the batch dimension broadcastable
# 🚀 Feature
Support different key/query/value batch sizes, if 1 or the same value

## Motivation
(popped up when working on [salina](https://github.com/facebookresearch/salina)) PyTorch MHA handles key/query/value with 1 to B batch sizes, will broadcast if in need, We don´t do that and hard crash instead, so the users have to align dimensions from the outside

## Pitch

<!-- A clear and concise description of what you want to happen. -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->


--- 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/multi_head_dispatch.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 dataclasses import asdict, dataclass
8 from typing import Optional
9 
10 import torch
11 import torch.nn as nn
12 from torch.nn.init import constant_
13 
14 from xformers.components.attention import Attention
15 from xformers.components.in_proj_container import InProjContainer, InProjParams
16 from xformers.components.positional_embedding import RotaryEmbedding
17 
18 
19 @dataclass
20 class MultiHeadDispatchConfig:
21     dim_model: int
22     residual_dropout: float
23     num_heads: int
24     attention: Attention
25     bias: bool
26     dim_key: Optional[int]
27     dim_value: Optional[int]
28     in_proj_container: Optional[InProjContainer]
29     use_separate_proj_weight: Optional[bool]
30     use_rotary_embeddings: Optional[bool]
31     out_proj: Optional[nn.Module]
32 
33 
34 # Move head forward and fold into batch dim. dimensions become (B * nh, S, hs)
35 def _fold_heads(t: torch.Tensor, B: int, S: int, H: int, Hs: int):
36     return t.view(B, S, H, Hs).transpose(1, 2).flatten(start_dim=0, end_dim=1)
37 
38 
39 def _split_heads(t: torch.Tensor, B: int, S: int, H: int, Hs: int):
40     return t.view(B, S, H, Hs).transpose(1, 2)
41 
42 
43 class MultiHeadDispatch(nn.Module):
44     """
45     A multi-head masked self-attention dispatch mechanism, with a projection at the end,
46     following the architecture proposed in `Attention is all you need`_, Vaswani et al.
47 
48     The actual attention mechanism can vary, as well as the projections.
49     This can be used to wrap the proposed attention mechanisms and make them multi-head aware,
50     but it is optional.
51 
52     .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762v5
53     """
54 
55     def __init__(
56         self,
57         dim_model: int,
58         residual_dropout: float,
59         num_heads: int,
60         attention: Attention,
61         bias: bool = True,
62         dim_key: Optional[int] = None,
63         dim_value: Optional[int] = None,
64         in_proj_container: Optional[InProjContainer] = None,
65         use_separate_proj_weight: Optional[bool] = False,
66         use_rotary_embeddings: Optional[bool] = False,
67         out_proj: Optional[nn.Module] = None,
68         *args,
69         **kwargs,
70     ):
71         super().__init__()
72 
73         assert (
74             dim_model % num_heads == 0
75         )  # static preset for now, each head works on 1/d the embeddings, could be relaxed
76         assert num_heads > 0
77 
78         # Popular default is that all latent dimensions are the same
79         dim_key, dim_value = map(lambda x: x if x else dim_model, (dim_key, dim_value))
80 
81         self.num_heads = num_heads
82         self.dim_k = dim_key // num_heads
83         self.dim_value = dim_value
84         self.dim_model = dim_model
85         self.attention = attention
86 
87         # key, query, value projections for all heads
88         # critical options are
89         # - are we sharing weights ?
90         # - are we adding biases, and if yes are they shared ?
91         if attention.requires_input_projection:
92             self.in_proj_container = (
93                 in_proj_container
94                 if in_proj_container is not None
95                 else InProjContainer(
96                     query_proj_params=InProjParams(dim_model, dim_key, bias=bias),
97                     key_proj_params=InProjParams(dim_model, dim_key, bias=bias)
98                     if use_separate_proj_weight
99                     else None,
100                     value_proj_params=InProjParams(dim_model, dim_value, bias=bias)
101                     if use_separate_proj_weight
102                     else None,
103                 )
104             )
105 
106         # Optional rotary embeddings
107         self.rotary_embeddings = (
108             RotaryEmbedding(self.dim_k) if use_rotary_embeddings else None
109         )
110 
111         # Regularization
112         self.resid_drop = nn.Dropout(residual_dropout, inplace=False)
113 
114         # Output projection
115         self.proj = out_proj if out_proj else nn.Linear(dim_model, dim_model, bias=bias)
116         if isinstance(self.proj, nn.Linear) and self.proj.bias is not None:
117             constant_(self.proj.bias, 0.0)
118 
119     def _check(self, t, name):
120         assert (
121             t.shape[2] % self.dim_k == 0
122         ), f"the {name} embeddings need to be divisible by the number of heads"
123 
124     def forward(
125         self,
126         query: torch.Tensor,
127         key: Optional[torch.Tensor] = None,
128         value: Optional[torch.Tensor] = None,
129         att_mask: Optional[torch.Tensor] = None,
130         key_padding_mask: Optional[torch.Tensor] = None,
131     ) -> torch.Tensor:
132         """
133         Expected input dimensions are [batch size, sequence length, embed dim]
134         Output dimensions are [batch size, sequence length, embed dim]
135         """
136 
137         if key is None:
138             key = query
139         if value is None:
140             value = query
141 
142         # Check the dimensions properly
143         self._check(query, "query")
144         self._check(value, "value")
145         self._check(key, "key")
146 
147         B, S_Q, _ = query.size()  # Batch x Sequence x Embedding (latent)
148         _, S_K, _ = key.size()  # K, Q's sequence length could differ
149 
150         # Catch different query and key length but a causal attention
151         if S_Q != S_K:
152             assert (
153                 not self.attention.requires_same_k_q_dimensions
154             ), "This attention mechanism requires query and key to have the same sequence (context) lengths"
155 
156             if hasattr(self.attention, "causal"):
157                 assert not self.attention.causal, (
158                     "Causal attention is not supported when key and query have different sequence lengths.\n"
159                     + "In that case causality is ill-determined. Please pad your sequences accordingly"
160                 )
161 
162         if self.attention.requires_skip_multi_head:
163             return self.attention(
164                 query, key, value, att_mask=att_mask, key_padding_mask=key_padding_mask
165             )
166 
167         # Calculate query, key, values for all heads in batch
168         if self.attention.requires_input_projection:
169             q, k, v = self.in_proj_container(query=query, key=key, value=value)
170         else:
171             k, q, v = key, query, value
172 
173         # Optional: rotary embedding, add relative positioning information
174         if self.rotary_embeddings:
175             # rotary requires the head dimension
176             q = _split_heads(q, B, S_Q, self.num_heads, self.dim_k)
177             k = _split_heads(k, B, S_K, self.num_heads, self.dim_k)
178             v = _split_heads(v, B, S_K, self.num_heads, self.dim_k)
179 
180             q, k = self.rotary_embeddings(q=q, k=k)
181 
182             if not self.attention.requires_head_dimension:
183                 q, k, v = q.flatten(0, 1), k.flatten(0, 1), v.flatten(0, 1)
184 
185         else:
186             # Reshape k/q/v to either expose the heads, or fold the head dimension into the batch
187             reshape_fn = (
188                 _split_heads if self.attention.requires_head_dimension else _fold_heads
189             )
190 
191             q = reshape_fn(q, B, S_Q, self.num_heads, self.dim_k)
192             k = reshape_fn(k, B, S_K, self.num_heads, self.dim_k)
193             v = reshape_fn(v, B, S_K, self.num_heads, self.dim_k)
194 
195         # Self-attend
196         y = self.attention(
197             q=q, k=k, v=v, att_mask=att_mask, key_padding_mask=key_padding_mask
198         )
199 
200         # Re-assemble all head outputs side by side
201         y = (
202             y.view(B, self.num_heads, S_Q, self.dim_k)
203             .transpose(1, 2)
204             .flatten(start_dim=2, end_dim=3)
205         )
206 
207         # Output projection, dropout and good to go
208         y = self.resid_drop(self.proj(y))
209 
210         # Return the same sequence size as the input
211         return y
212 
213     @classmethod
214     def from_config(cls, config: MultiHeadDispatchConfig):
215         # Generate the class inputs from the config
216         fields = asdict(config)
217 
218         # Skip all Nones so that default values are used
219         fields = {k: v for k, v in fields.items() if v is not None}
220 
221         return cls(**fields)
222
```
--- 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/multi_head_dispatch.py b/xformers/components/multi_head_dispatch.py
--- a/xformers/components/multi_head_dispatch.py
+++ b/xformers/components/multi_head_dispatch.py
@@ -143,6 +143,11 @@
         self._check(value, "value")
         self._check(key, "key")
 
+        max_batch = max((query.shape[0], key.shape[0], value.shape[0]))
+        query, key, value = map(
+            lambda x: x.expand(max_batch, -1, -1), [query, key, value]
+        )
+
         B, S_Q, _ = query.size()  # Batch x Sequence x Embedding (latent)
         _, S_K, _ = key.size()  # K, Q's sequence length could differ

{"golden_diff": "diff --git a/xformers/components/multi_head_dispatch.py b/xformers/components/multi_head_dispatch.py\n--- a/xformers/components/multi_head_dispatch.py\n+++ b/xformers/components/multi_head_dispatch.py\n@@ -143,6 +143,11 @@\n         self._check(value, \"value\")\n         self._check(key, \"key\")\n \n+        max_batch = max((query.shape[0], key.shape[0], value.shape[0]))\n+        query, key, value = map(\n+            lambda x: x.expand(max_batch, -1, -1), [query, key, value]\n+        )\n+\n         B, S_Q, _ = query.size()  # Batch x Sequence x Embedding (latent)\n         _, S_K, _ = key.size()  # K, Q's sequence length could differ\n", "issue": "[feat] Make the batch dimension broadcastable\n# \ud83d\ude80 Feature\r\nSupport different key/query/value batch sizes, if 1 or the same value\r\n\r\n## Motivation\r\n(popped up when working on [salina](https://github.com/facebookresearch/salina)) PyTorch MHA handles key/query/value with 1 to B batch sizes, will broadcast if in need, We don\u00b4t do that and hard crash instead, so the users have to align dimensions from the outside\r\n\r\n## Pitch\r\n\r\n<!-- A clear and concise description of what you want to happen. -->\r\n\r\n## Alternatives\r\n\r\n<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->\r\n\r\n## Additional context\r\n\r\n<!-- Add any other context or screenshots about the feature request here. -->\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 dataclasses import asdict, dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn.init import constant_\n\nfrom xformers.components.attention import Attention\nfrom xformers.components.in_proj_container import InProjContainer, InProjParams\nfrom xformers.components.positional_embedding import RotaryEmbedding\n\n\n@dataclass\nclass MultiHeadDispatchConfig:\n    dim_model: int\n    residual_dropout: float\n    num_heads: int\n    attention: Attention\n    bias: bool\n    dim_key: Optional[int]\n    dim_value: Optional[int]\n    in_proj_container: Optional[InProjContainer]\n    use_separate_proj_weight: Optional[bool]\n    use_rotary_embeddings: Optional[bool]\n    out_proj: Optional[nn.Module]\n\n\n# Move head forward and fold into batch dim. dimensions become (B * nh, S, hs)\ndef _fold_heads(t: torch.Tensor, B: int, S: int, H: int, Hs: int):\n    return t.view(B, S, H, Hs).transpose(1, 2).flatten(start_dim=0, end_dim=1)\n\n\ndef _split_heads(t: torch.Tensor, B: int, S: int, H: int, Hs: int):\n    return t.view(B, S, H, Hs).transpose(1, 2)\n\n\nclass MultiHeadDispatch(nn.Module):\n    \"\"\"\n    A multi-head masked self-attention dispatch mechanism, with a projection at the end,\n    following the architecture proposed in `Attention is all you need`_, Vaswani et al.\n\n    The actual attention mechanism can vary, as well as the projections.\n    This can be used to wrap the proposed attention mechanisms and make them multi-head aware,\n    but it is optional.\n\n    .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762v5\n    \"\"\"\n\n    def __init__(\n        self,\n        dim_model: int,\n        residual_dropout: float,\n        num_heads: int,\n        attention: Attention,\n        bias: bool = True,\n        dim_key: Optional[int] = None,\n        dim_value: Optional[int] = None,\n        in_proj_container: Optional[InProjContainer] = None,\n        use_separate_proj_weight: Optional[bool] = False,\n        use_rotary_embeddings: Optional[bool] = False,\n        out_proj: Optional[nn.Module] = None,\n        *args,\n        **kwargs,\n    ):\n        super().__init__()\n\n        assert (\n            dim_model % num_heads == 0\n        )  # static preset for now, each head works on 1/d the embeddings, could be relaxed\n        assert num_heads > 0\n\n        # Popular default is that all latent dimensions are the same\n        dim_key, dim_value = map(lambda x: x if x else dim_model, (dim_key, dim_value))\n\n        self.num_heads = num_heads\n        self.dim_k = dim_key // num_heads\n        self.dim_value = dim_value\n        self.dim_model = dim_model\n        self.attention = attention\n\n        # key, query, value projections for all heads\n        # critical options are\n        # - are we sharing weights ?\n        # - are we adding biases, and if yes are they shared ?\n        if attention.requires_input_projection:\n            self.in_proj_container = (\n                in_proj_container\n                if in_proj_container is not None\n                else InProjContainer(\n                    query_proj_params=InProjParams(dim_model, dim_key, bias=bias),\n                    key_proj_params=InProjParams(dim_model, dim_key, bias=bias)\n                    if use_separate_proj_weight\n                    else None,\n                    value_proj_params=InProjParams(dim_model, dim_value, bias=bias)\n                    if use_separate_proj_weight\n                    else None,\n                )\n            )\n\n        # Optional rotary embeddings\n        self.rotary_embeddings = (\n            RotaryEmbedding(self.dim_k) if use_rotary_embeddings else None\n        )\n\n        # Regularization\n        self.resid_drop = nn.Dropout(residual_dropout, inplace=False)\n\n        # Output projection\n        self.proj = out_proj if out_proj else nn.Linear(dim_model, dim_model, bias=bias)\n        if isinstance(self.proj, nn.Linear) and self.proj.bias is not None:\n            constant_(self.proj.bias, 0.0)\n\n    def _check(self, t, name):\n        assert (\n            t.shape[2] % self.dim_k == 0\n        ), f\"the {name} embeddings need to be divisible by the number of heads\"\n\n    def forward(\n        self,\n        query: torch.Tensor,\n        key: Optional[torch.Tensor] = None,\n        value: Optional[torch.Tensor] = None,\n        att_mask: Optional[torch.Tensor] = None,\n        key_padding_mask: Optional[torch.Tensor] = None,\n    ) -> torch.Tensor:\n        \"\"\"\n        Expected input dimensions are [batch size, sequence length, embed dim]\n        Output dimensions are [batch size, sequence length, embed dim]\n        \"\"\"\n\n        if key is None:\n            key = query\n        if value is None:\n            value = query\n\n        # Check the dimensions properly\n        self._check(query, \"query\")\n        self._check(value, \"value\")\n        self._check(key, \"key\")\n\n        B, S_Q, _ = query.size()  # Batch x Sequence x Embedding (latent)\n        _, S_K, _ = key.size()  # K, Q's sequence length could differ\n\n        # Catch different query and key length but a causal attention\n        if S_Q != S_K:\n            assert (\n                not self.attention.requires_same_k_q_dimensions\n            ), \"This attention mechanism requires query and key to have the same sequence (context) lengths\"\n\n            if hasattr(self.attention, \"causal\"):\n                assert not self.attention.causal, (\n                    \"Causal attention is not supported when key and query have different sequence lengths.\\n\"\n                    + \"In that case causality is ill-determined. Please pad your sequences accordingly\"\n                )\n\n        if self.attention.requires_skip_multi_head:\n            return self.attention(\n                query, key, value, att_mask=att_mask, key_padding_mask=key_padding_mask\n            )\n\n        # Calculate query, key, values for all heads in batch\n        if self.attention.requires_input_projection:\n            q, k, v = self.in_proj_container(query=query, key=key, value=value)\n        else:\n            k, q, v = key, query, value\n\n        # Optional: rotary embedding, add relative positioning information\n        if self.rotary_embeddings:\n            # rotary requires the head dimension\n            q = _split_heads(q, B, S_Q, self.num_heads, self.dim_k)\n            k = _split_heads(k, B, S_K, self.num_heads, self.dim_k)\n            v = _split_heads(v, B, S_K, self.num_heads, self.dim_k)\n\n            q, k = self.rotary_embeddings(q=q, k=k)\n\n            if not self.attention.requires_head_dimension:\n                q, k, v = q.flatten(0, 1), k.flatten(0, 1), v.flatten(0, 1)\n\n        else:\n            # Reshape k/q/v to either expose the heads, or fold the head dimension into the batch\n            reshape_fn = (\n                _split_heads if self.attention.requires_head_dimension else _fold_heads\n            )\n\n            q = reshape_fn(q, B, S_Q, self.num_heads, self.dim_k)\n            k = reshape_fn(k, B, S_K, self.num_heads, self.dim_k)\n            v = reshape_fn(v, B, S_K, self.num_heads, self.dim_k)\n\n        # Self-attend\n        y = self.attention(\n            q=q, k=k, v=v, att_mask=att_mask, key_padding_mask=key_padding_mask\n        )\n\n        # Re-assemble all head outputs side by side\n        y = (\n            y.view(B, self.num_heads, S_Q, self.dim_k)\n            .transpose(1, 2)\n            .flatten(start_dim=2, end_dim=3)\n        )\n\n        # Output projection, dropout and good to go\n        y = self.resid_drop(self.proj(y))\n\n        # Return the same sequence size as the input\n        return y\n\n    @classmethod\n    def from_config(cls, config: MultiHeadDispatchConfig):\n        # Generate the class inputs from the config\n        fields = asdict(config)\n\n        # Skip all Nones so that default values are used\n        fields = {k: v for k, v in fields.items() if v is not None}\n\n        return cls(**fields)\n", "path": "xformers/components/multi_head_dispatch.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 dataclasses import asdict, dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn.init import constant_\n\nfrom xformers.components.attention import Attention\nfrom xformers.components.in_proj_container import InProjContainer, InProjParams\nfrom xformers.components.positional_embedding import RotaryEmbedding\n\n\n@dataclass\nclass MultiHeadDispatchConfig:\n    dim_model: int\n    residual_dropout: float\n    num_heads: int\n    attention: Attention\n    bias: bool\n    dim_key: Optional[int]\n    dim_value: Optional[int]\n    in_proj_container: Optional[InProjContainer]\n    use_separate_proj_weight: Optional[bool]\n    use_rotary_embeddings: Optional[bool]\n    out_proj: Optional[nn.Module]\n\n\n# Move head forward and fold into batch dim. dimensions become (B * nh, S, hs)\ndef _fold_heads(t: torch.Tensor, B: int, S: int, H: int, Hs: int):\n    return t.view(B, S, H, Hs).transpose(1, 2).flatten(start_dim=0, end_dim=1)\n\n\ndef _split_heads(t: torch.Tensor, B: int, S: int, H: int, Hs: int):\n    return t.view(B, S, H, Hs).transpose(1, 2)\n\n\nclass MultiHeadDispatch(nn.Module):\n    \"\"\"\n    A multi-head masked self-attention dispatch mechanism, with a projection at the end,\n    following the architecture proposed in `Attention is all you need`_, Vaswani et al.\n\n    The actual attention mechanism can vary, as well as the projections.\n    This can be used to wrap the proposed attention mechanisms and make them multi-head aware,\n    but it is optional.\n\n    .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762v5\n    \"\"\"\n\n    def __init__(\n        self,\n        dim_model: int,\n        residual_dropout: float,\n        num_heads: int,\n        attention: Attention,\n        bias: bool = True,\n        dim_key: Optional[int] = None,\n        dim_value: Optional[int] = None,\n        in_proj_container: Optional[InProjContainer] = None,\n        use_separate_proj_weight: Optional[bool] = False,\n        use_rotary_embeddings: Optional[bool] = False,\n        out_proj: Optional[nn.Module] = None,\n        *args,\n        **kwargs,\n    ):\n        super().__init__()\n\n        assert (\n            dim_model % num_heads == 0\n        )  # static preset for now, each head works on 1/d the embeddings, could be relaxed\n        assert num_heads > 0\n\n        # Popular default is that all latent dimensions are the same\n        dim_key, dim_value = map(lambda x: x if x else dim_model, (dim_key, dim_value))\n\n        self.num_heads = num_heads\n        self.dim_k = dim_key // num_heads\n        self.dim_value = dim_value\n        self.dim_model = dim_model\n        self.attention = attention\n\n        # key, query, value projections for all heads\n        # critical options are\n        # - are we sharing weights ?\n        # - are we adding biases, and if yes are they shared ?\n        if attention.requires_input_projection:\n            self.in_proj_container = (\n                in_proj_container\n                if in_proj_container is not None\n                else InProjContainer(\n                    query_proj_params=InProjParams(dim_model, dim_key, bias=bias),\n                    key_proj_params=InProjParams(dim_model, dim_key, bias=bias)\n                    if use_separate_proj_weight\n                    else None,\n                    value_proj_params=InProjParams(dim_model, dim_value, bias=bias)\n                    if use_separate_proj_weight\n                    else None,\n                )\n            )\n\n        # Optional rotary embeddings\n        self.rotary_embeddings = (\n            RotaryEmbedding(self.dim_k) if use_rotary_embeddings else None\n        )\n\n        # Regularization\n        self.resid_drop = nn.Dropout(residual_dropout, inplace=False)\n\n        # Output projection\n        self.proj = out_proj if out_proj else nn.Linear(dim_model, dim_model, bias=bias)\n        if isinstance(self.proj, nn.Linear) and self.proj.bias is not None:\n            constant_(self.proj.bias, 0.0)\n\n    def _check(self, t, name):\n        assert (\n            t.shape[2] % self.dim_k == 0\n        ), f\"the {name} embeddings need to be divisible by the number of heads\"\n\n    def forward(\n        self,\n        query: torch.Tensor,\n        key: Optional[torch.Tensor] = None,\n        value: Optional[torch.Tensor] = None,\n        att_mask: Optional[torch.Tensor] = None,\n    ) -> torch.Tensor:\n        \"\"\"\n        Expected input dimensions are [batch size, sequence length, embed dim]\n        Output dimensions are [batch size, sequence length, embed dim]\n        \"\"\"\n\n        if key is None:\n            key = query\n        if value is None:\n            value = query\n\n        # Check the dimensions properly\n        self._check(query, \"query\")\n        self._check(value, \"value\")\n        self._check(key, \"key\")\n\n        max_batch = max((query.shape[0], key.shape[0], value.shape[0]))\n        query, key, value = map(\n            lambda x: x.expand(max_batch, -1, -1), [query, key, value]\n        )\n\n        B, S_Q, _ = query.size()  # Batch x Sequence x Embedding (latent)\n        _, S_K, _ = key.size()  # K, Q's sequence length could differ\n\n        # Catch different query and key length but a causal attention\n        if S_Q != S_K:\n            assert (\n                not self.attention.requires_same_k_q_dimensions\n            ), \"This attention mechanism requires query and key to have the same sequence (context) lengths\"\n\n            if hasattr(self.attention, \"causal\"):\n                assert not self.attention.causal, (\n                    \"Causal attention is not supported when key and query have different sequence lengths.\\n\"\n                    + \"In that case causality is ill-determined. Please pad your sequences accordingly\"\n                )\n\n        # Calculate query, key, values for all heads in batch\n        if self.attention.requires_input_projection:\n            q, k, v = self.in_proj_container(query=query, key=key, value=value)\n        else:\n            k, q, v = key, query, value\n\n        # Optional: rotary embedding, add relative positioning information\n        if self.rotary_embeddings:\n            # rotary requires the head dimension\n            q = _split_heads(q, B, S_Q, self.num_heads, self.dim_k)\n            k = _split_heads(k, B, S_K, self.num_heads, self.dim_k)\n            v = _split_heads(v, B, S_K, self.num_heads, self.dim_k)\n\n            q, k = self.rotary_embeddings(q=q, k=k)\n\n            if not self.attention.requires_head_dimension:\n                q, k, v = q.flatten(0, 1), k.flatten(0, 1), v.flatten(0, 1)\n\n        else:\n            # Reshape k/q/v to either expose the heads, or fold the head dimension into the batch\n            reshape_fn = (\n                _split_heads if self.attention.requires_head_dimension else _fold_heads\n            )\n\n            q = reshape_fn(q, B, S_Q, self.num_heads, self.dim_k)\n            k = reshape_fn(k, B, S_K, self.num_heads, self.dim_k)\n            v = reshape_fn(v, B, S_K, self.num_heads, self.dim_k)\n\n        # Self-attend\n        y = self.attention(q=q, k=k, v=v, att_mask=att_mask)\n\n        # Re-assemble all head outputs side by side\n        y = (\n            y.view(B, self.num_heads, S_Q, self.dim_k)\n            .transpose(1, 2)\n            .flatten(start_dim=2, end_dim=3)\n        )\n\n        # Output projection, dropout and good to go\n        y = self.resid_drop(self.proj(y))\n\n        # Return the same sequence size as the input\n        return y\n\n    @classmethod\n    def from_config(cls, config: MultiHeadDispatchConfig):\n        # Generate the class inputs from the config\n        fields = asdict(config)\n\n        # Skip all Nones so that default values are used\n        fields = {k: v for k, v in fields.items() if v is not None}\n\n        return cls(**fields)\n", "path": "xformers/components/multi_head_dispatch.py"}]}