x-transformers 2.11.7__py3-none-any.whl → 2.11.8__py3-none-any.whl

x_transformers/free_transformer.py

@@ -296,10 +296,7 @@ class FreeTransformer(Module):
 
             head_embed = self.decoder_head(tokens)
 
-            if exists(condition):
-                head_embed = head_embed + condition
-
-            tail_embed = self.decoder_tail(head_embed)
+            tail_embed = self.decoder_tail(head_embed, self_attn_kv_residuals = condition)
 
             tail_embed = tail_embed[:, -1]
 
@@ -339,7 +336,7 @@ class FreeTransformer(Module):
 
         # decoder tail
 
-        tokens = self.decoder_tail(tokens + condition)
+        tokens = self.decoder_tail(tokens, self_attn_kv_residuals = condition)
 
         # cross entropy loss
 

x_transformers/x_transformers.py

@@ -1686,6 +1686,7 @@ class Attention(Module):
         value_residual = None,
         additional_key_values: tuple[Tensor, Tensor] | None = None,
         additional_key_value_mask = None,
+        kv_input_residual = 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
 
@@ -1702,6 +1703,12 @@ class Attention(Module):
             kv_input = default(context, x)
             q_input, k_input, v_input = x, kv_input, kv_input
 
+        # done for free transformer
+
+        if exists(kv_input_residual):
+            k_input = k_input + kv_input_residual
+            v_input = v_input + kv_input_residual
+
         if exists(mem):
             k_input, mem_packed_shape = pack([mem, k_input], 'b * d')
             v_input, _ = pack([mem, v_input], 'b * d')
@@ -2543,7 +2550,9 @@ class AttentionLayers(Module):
         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
+        layers_execute_order: tuple[int, ...] | None = None,
+        self_attn_kv_residuals: Tensor | None = None,
+        cross_attn_kv_residuals: Tensor | None = None
     ):
         assert not (self.cross_attend ^ exists(context)), 'context must be passed in if cross_attend is set to True'
         assert not (exists(condition) ^ self.need_condition), 'condition needs to be passed in if using adaptive layernorm or vice versa'
@@ -2721,6 +2730,23 @@ class AttentionLayers(Module):
 
         skip_hiddens = []
 
+        # for residuals to key value inputs for self and cross attention
+
+        self_attn_kv_residuals_iter = iter((None,))
+        cross_attn_kv_residuals_iter = iter((None,))
+
+        if exists(self_attn_kv_residuals):
+            if self_attn_kv_residuals.ndim == 3:
+                self_attn_kv_residuals = rearrange(self_attn_kv_residuals, '... ->  1 ...')
+
+            self_attn_kv_residuals_iter = iter(self_attn_kv_residuals)
+
+        if exists(cross_attn_kv_residuals):
+            if cross_attn_kv_residuals.ndim == 3:
+                cross_attn_kv_residuals = rearrange(cross_attn_kv_residuals, '... ->  1 ...')
+
+            cross_attn_kv_residuals_iter = iter(cross_attn_kv_residuals)
+
         # for value residuals
 
         first_self_attn_inter = None
@@ -2794,9 +2820,9 @@ 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, 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)
+                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, kv_input_residual = next(self_attn_kv_residuals_iter, None), 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)
+                out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), kv_input_residual = next(cross_attn_kv_residuals_iter, None), value_residual = maybe_cross_attn_value_residual, **cross_attn_rotary_pos_emb, return_intermediates = True)
             elif layer_type == 'f':
                 out = block(x, deep_embed = next(deep_embeds_iter, None))
 

{x_transformers-2.11.7.dist-info → x_transformers-2.11.8.dist-info}/METADATA

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

{x_transformers-2.11.7.dist-info → x_transformers-2.11.8.dist-info}/RECORD

@@ -5,16 +5,16 @@ x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTN
 x_transformers/continuous.py,sha256=WwpQCjyVY4PtuEAOFY68zqgklbF9I7AL5w6874YlDe8,13249
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
-x_transformers/free_transformer.py,sha256=RVEMvmiSVrTJeLbXVoI4YeR26kP2mlqmwFHHLw7IJ_Q,9582
+x_transformers/free_transformer.py,sha256=kfl_MIZxv4TARRQbq3NroGwZSBVHdYoJNu1hfWMloco,9555
 x_transformers/gpt_vae.py,sha256=4QdznXZcU7pmMXUeEocAOKpcTkREYS-zDHktN5ADtNk,5981
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=ADr83Fz2cehj_F7N1bMwxhAg-r48fGhlaZqw3hxoxMQ,125765
+x_transformers/x_transformers.py,sha256=bYnVtkcfr082ALprIGgYIUx53lLADGYpi9t6QEJp1Kc,126907
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.11.7.dist-info/METADATA,sha256=1rStN6vwvwU-0hm17zWx8yXL99ZrZ-_KGMsFwvRIPVg,96011
-x_transformers-2.11.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.11.7.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.11.7.dist-info/RECORD,,
+x_transformers-2.11.8.dist-info/METADATA,sha256=NTTtQVh5bRCnk7RDpma7JHairCWIvaO2euEE2djXUFA,96011
+x_transformers-2.11.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.8.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.8.dist-info/RECORD,,