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

x_transformers/x_transformers.py

@@ -1795,6 +1795,13 @@ class Attention(Module):
             seq_len = k.shape[-2]
 
             added_k, added_v = additional_key_values
+            added_kv_heads, added_kv_len = added_k.shape[1], added_k.shape[-2]
+
+            # take care of expanding to query heads if mismatch between key / value heads with the ones coming from vlm
+
+            if added_kv_heads != kv_h:
+                assert divisible_by(h, added_kv_heads)
+                k, v, added_k, added_v = tuple(repeat(t, 'b h ... -> b (r h) ...', r = h // t.shape[1]) for t in (k, v, added_k, added_v))
 
             k = cat((added_k, k), dim = -2)
             v = cat((added_v, v), dim = -2)
@@ -1802,7 +1809,6 @@ class Attention(Module):
             if (exists(input_mask) or exists(additional_key_value_mask)):
 
                 if not exists(additional_key_value_mask):
-                    added_kv_len = added_k.shape[-2]
                     input_mask = pad_at_dim(input_mask, (added_kv_len, 0), dim = -1, value = True)
                 elif not exists(input_mask):
                     input_mask = pad_at_dim(additional_key_value_mask, (0, seq_len), dim = -1, value = True)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.6.2
+Version: 2.6.3
 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.2.dist-info → x_transformers-2.6.3.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=kcdU6gp4QXfab9P0M5WYnVic6nTFrTtRGyLEcPFBQcY,121719
+x_transformers/x_transformers.py,sha256=B7dv_LuzODwCrTsfDnp28g-_lMnirQE3gteQwSGyW5k,122100
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-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,,
+x_transformers-2.6.3.dist-info/METADATA,sha256=DaTrEChlXc_zpUXv-Jw3A4ca4aon0Ons7wl4-wj1XzY,90223
+x_transformers-2.6.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.6.3.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.6.3.dist-info/RECORD,,