x-transformers 2.11.8__tar.gz → 2.11.10__tar.gz

{x_transformers-2.11.8 → x_transformers-2.11.10}/PKG-INFO

@@ -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 → x_transformers-2.11.10}/pyproject.toml

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

{x_transformers-2.11.8 → x_transformers-2.11.10}/tests/test_x_transformers.py

@@ -1410,7 +1410,9 @@ def test_attn_negative_weights(
     logits = model(x)
 
 @param('per_token_latents', (False, True))
+@param('dec_head_depth', (0, 4))
 def test_free(
+    dec_head_depth,
     per_token_latents
 ):
     from x_transformers.free_transformer import FreeTransformer
@@ -1420,9 +1422,9 @@ def test_free(
         max_seq_len = 1024,
         dim = 512,
         heads = 8,
-        dec_head_depth = 4,
+        dec_head_depth = dec_head_depth,
         dec_tail_depth = 4,
-        enc_depth = 3,
+        enc_depth = 2,
         kl_loss_weight = 1.,
         per_token_latents = per_token_latents,
         latent_bits = 8

{x_transformers-2.11.8 → x_transformers-2.11.10}/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