x-transformers 2.6.1__tar.gz → 2.6.3__tar.gz

{x_transformers-2.6.1 → x_transformers-2.6.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.6.1
+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.1 → x_transformers-2.6.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.6.1"
+version = "2.6.3"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.6.1 → x_transformers-2.6.3}/tests/test_x_transformers.py

@@ -1219,7 +1219,7 @@ def test_external_key_values():
         max_seq_len = 1024,
         attn_layers = Decoder(
             dim = 512,
-            depth = 2,
+            depth = 3,
             heads = 8,
             attn_dim_head = 16
         )
@@ -1228,8 +1228,8 @@ def test_external_key_values():
     seq = torch.randint(0, 20000, (3, 1024))
 
     key_values = [
-        (torch.randn(3, 8, 32, 16), torch.randn(3, 8, 32, 16)),
-        (torch.randn(3, 8, 32, 16), torch.randn(3, 8, 32, 16)),
+        (torch.randn(3, 2, 32, 16), torch.randn(3, 2, 32, 16)),
+        (torch.randn(3, 2, 32, 16), torch.randn(3, 2, 32, 16)),
     ]
 
     additional_kv_mask = torch.randint(0, 2, (3, 32)).bool()

{x_transformers-2.6.1 → x_transformers-2.6.3}/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)
@@ -2435,6 +2441,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 +2551,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 +2585,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