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

x_transformers/free_transformer.py

@@ -197,7 +197,9 @@ class FreeTransformer(Module):
             pre_norm_has_final_norm = False,
             **kwargs,
             **dec_kwargs
-        )
+        ) if dec_head_depth > 0 else None
+
+        assert dec_tail_depth > 0
 
         self.decoder_tail = Decoder(
             dim = dim,
@@ -266,7 +268,8 @@ class FreeTransformer(Module):
         seq_len,
         latents = None,
         filter_logits_fn = top_p,
-        logit_filter_kwargs: dict = dict(thres = 0.9)
+        logit_filter_kwargs: dict = dict(thres = 0.9),
+        use_kv_cache = True
     ):
         prompts, inverse_pack = pack_with_inverse(prompts, '* n')
 
@@ -280,10 +283,16 @@ class FreeTransformer(Module):
                 latents = tensor(latents, device = self.device)
 
             if latents.ndim == 1: # repeat latents
-                latents = repeat(latents, 'd -> b d', b = batch)
+                latents = repeat(latents, 'd -> b 1 d', b = batch)
+            elif latents.ndim == 2:
+                latents = rearrange(latents, 'b d -> b 1 d')
 
             condition = self.from_latent_to_condition(latents)
 
+        # kv cache
+
+        head_cache = tail_cache = None
+
         # generated
 
         prompt_len = prompts.shape[-1]
@@ -294,9 +303,20 @@ class FreeTransformer(Module):
 
         for _ in range(max(0, seq_len - prompt_len)):
 
-            head_embed = self.decoder_head(tokens)
+            # head, which may not exist
+
+            if exists(self.decoder_head):
+                head_embed, next_head_cache = self.decoder_head(tokens, cache = head_cache, return_hiddens = True)
+            else:
+                head_embed, next_head_cache = tokens, None
 
-            tail_embed = self.decoder_tail(head_embed, self_attn_kv_residuals = condition)
+            # handle one token being given to the decoder tail when doing kv caching - rotary embedding needs to know the seq position offset
+
+            seq_pos_offset = head_cache.cache_length if exists(head_cache) else 0
+
+            # tail
+
+            tail_embed, next_tail_cache = self.decoder_tail(head_embed, cache = tail_cache, seq_pos_offset = seq_pos_offset, self_attn_kv_residuals = condition, return_hiddens = True)
 
             tail_embed = tail_embed[:, -1]
 
@@ -309,6 +329,10 @@ class FreeTransformer(Module):
             generated, _ = pack((generated, sampled), 'b *')
             tokens, _ = pack((tokens, self.token_emb(sampled)), 'b * d')
 
+            if use_kv_cache:
+                head_cache = next_head_cache
+                tail_cache = next_tail_cache
+
         return inverse_pack(generated)
 
     def forward(
@@ -326,7 +350,8 @@ class FreeTransformer(Module):
 
         # decoder head
 
-        tokens = self.decoder_head(tokens)
+        if exists(self.decoder_head):
+            tokens = self.decoder_head(tokens)
 
         # get latent Z
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.8
+Version: 2.11.10
 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.8.dist-info → x_transformers-2.11.10.dist-info}/RECORD

@@ -5,7 +5,7 @@ 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=kfl_MIZxv4TARRQbq3NroGwZSBVHdYoJNu1hfWMloco,9555
+x_transformers/free_transformer.py,sha256=9vpZLTy1hWrX_wZpxmCsoJjsoKisu7HO4MTzP48oZoQ,10535
 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
@@ -14,7 +14,7 @@ x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,
 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.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,,
+x_transformers-2.11.10.dist-info/METADATA,sha256=xcHidmoWV-DKOo65NAd84GKA4kRQwSGMvWh69Rwh_w8,96012
+x_transformers-2.11.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.10.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.10.dist-info/RECORD,,