x-transformers 2.6.1__py3-none-any.whl → 2.6.2__py3-none-any.whl

x_transformers/x_transformers.py

@@ -2435,6 +2435,7 @@ class AttentionLayers(Module):
         deep_embeds_and_ids: tuple[nn.Parameter, Tensor] | None = None,
         self_attn_additional_kv: list[tuple[Tensor, Tensor]] | None = None,
         additional_kv_mask = None,
+        route_additional_kv_to_top = True,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
         layers_execute_order: tuple[int, ...] | None = None
@@ -2544,10 +2545,6 @@ class AttentionLayers(Module):
 
         iter_attn_cache = iter(attn_cache)
 
-        # additional self attn key / values
-
-        iter_self_attn_kv = iter(default(self_attn_additional_kv, ()))
-
         # handle deep embeds if needed
 
         deep_embeds = []
@@ -2582,6 +2579,16 @@ class AttentionLayers(Module):
         layers_execute_order = default(layers_execute_order, self.layers_execute_order)
         layer_variables = tuple(tuple(layer_variable[i] for i in layers_execute_order) for layer_variable in layer_variables)
 
+        # additional self attn key / values - say coming from vlm
+
+        if exists(self_attn_additional_kv) and route_additional_kv_to_top:
+            num_self_attns = sum([layer_type == 'a' for layer_type in first(layer_variables)])
+
+            self_attn_additional_kv = self_attn_additional_kv[-num_self_attns:]
+            self_attn_additional_kv = [None] * (num_self_attns - len(self_attn_additional_kv)) + self_attn_additional_kv
+
+        iter_self_attn_kv = iter(default(self_attn_additional_kv, ()))
+
         # derived input for reinjection if needed
 
         inp_inject = None

{x_transformers-2.6.1.dist-info → x_transformers-2.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.6.1
+Version: 2.6.2
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.6.1.dist-info → x_transformers-2.6.2.dist-info}/RECORD

@@ -9,10 +9,10 @@ x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=O8Z4j7wDrj47-lZxmpvToHbHpoFqLy2pk199tQ4v4hI,121281
+x_transformers/x_transformers.py,sha256=kcdU6gp4QXfab9P0M5WYnVic6nTFrTtRGyLEcPFBQcY,121719
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.6.1.dist-info/METADATA,sha256=VVzitcHytmh6tNmtjSMyDWjFxfjpQu6PhR4sTFkxjpk,90223
-x_transformers-2.6.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.6.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.6.1.dist-info/RECORD,,
+x_transformers-2.6.2.dist-info/METADATA,sha256=zHbOLX82fzv9Lzg6i2ZqkJhGJnDtAvwNIl9fbnKnhl8,90223
+x_transformers-2.6.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.6.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.6.2.dist-info/RECORD,,