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

x_transformers/x_transformers.py

@@ -1617,7 +1617,9 @@ class Attention(Module):
         mem_mask = None,
         return_intermediates = False,
         cache: Intermediates | None = None,
-        value_residual = None
+        value_residual = None,
+        additional_key_values: tuple[Tensor, Tensor] | None = None,
+        additional_key_value_mask = None,
     ):
         b, n, h, kv_h, head_scale, num_mem_kv, device, has_context, qkv_receive_diff_residuals, is_multi_latent_attn = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, self.num_mem_kv, x.device, exists(context), self.qkv_receive_diff_residuals, self.use_latent_kv
 
@@ -1787,6 +1789,26 @@ class Attention(Module):
             if exists(input_mask):
                 input_mask = pad_at_dim(input_mask, (self.num_mem_kv, 0), dim = -1, value = True)
 
+        # maybe append additional key / values
+
+        if exists(additional_key_values):
+            seq_len = k.shape[-2]
+
+            added_k, added_v = additional_key_values
+
+            k = cat((added_k, k), dim = -2)
+            v = cat((added_v, v), dim = -2)
+
+            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)
+                else:
+                    input_mask = cat((additional_key_value_mask, input_mask), dim = -1)
+
         # determine masking
 
         mask_value = max_neg_value(q)
@@ -2411,6 +2433,8 @@ class AttentionLayers(Module):
         context_pos = None,
         attn_bias = None,
         deep_embeds_and_ids: tuple[nn.Parameter, Tensor] | None = None,
+        self_attn_additional_kv: list[tuple[Tensor, Tensor]] | None = None,
+        additional_kv_mask = None,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
         layers_execute_order: tuple[int, ...] | None = None
@@ -2520,6 +2544,10 @@ 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 = []
@@ -2647,7 +2675,7 @@ class AttentionLayers(Module):
             # forward depending on layer type
 
             if layer_type == 'a':
-                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, value_residual = maybe_self_attn_value_residual, return_intermediates = True)
+                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, additional_key_values = next(iter_self_attn_kv, None), additional_key_value_mask = additional_kv_mask, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, value_residual = maybe_self_attn_value_residual, return_intermediates = True)
             elif layer_type == 'c':
                 out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), value_residual = maybe_cross_attn_value_residual, **cross_attn_rotary_pos_emb, return_intermediates = True)
             elif layer_type == 'f':

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

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

{x_transformers-2.5.6.dist-info → x_transformers-2.6.1.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=ymq4TL2OyeCPxwbTL0ShKptzRsMWXxsPv1MKZ9MxbHY,120021
+x_transformers/x_transformers.py,sha256=O8Z4j7wDrj47-lZxmpvToHbHpoFqLy2pk199tQ4v4hI,121281
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.5.6.dist-info/METADATA,sha256=2EpDDqrjJ20hyMPfHwQeoigq4LqJhD93h9w12FX82uQ,90223
-x_transformers-2.5.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.5.6.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.5.6.dist-info/RECORD,,
+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,,