x-transformers 2.3.1__tar.gz → 2.3.3__tar.gz

{x_transformers-2.3.1 → x_transformers-2.3.3}/PKG-INFO

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

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

{x_transformers-2.3.1 → x_transformers-2.3.3}/tests/test_x_transformers.py

@@ -849,8 +849,10 @@ def test_custom_ff_activation():
     assert logits.shape == (2, 1024, 20000)
 
 @pytest.mark.parametrize('probabilistic', (False, True))
+@pytest.mark.parametrize('cache_kv', (False, True))
 def test_continuous(
-    probabilistic
+    probabilistic,
+    cache_kv
 ):
     from x_transformers import (
         ContinuousTransformerWrapper,
@@ -887,5 +889,5 @@ def test_continuous(
     # then generate
 
     start_emb = torch.randn(1, 777)
-    generated = model.generate(start_emb, 17) # (17, 777)
+    generated = model.generate(start_emb, 17, cache_kv = cache_kv) # (17, 777)
     assert generated.shape == (17, 777)

{x_transformers-2.3.1 → x_transformers-2.3.3}/x_transformers/continuous.py

@@ -10,6 +10,7 @@ import einx
 from einops import rearrange, reduce, pack, repeat, unpack
 
 from x_transformers.x_transformers import (
+    Attention,
     AttentionLayers,
     ScaledSinusoidalEmbedding,
     AbsolutePositionalEmbedding,
@@ -111,6 +112,10 @@ class ContinuousTransformerWrapper(Module):
 
         self.project_out = nn.Linear(dim, dim_out * (2 if probabilistic else 1), bias = False) if exists(dim_out) else nn.Identity()
 
+        # can cache kv
+
+        self.can_cache_kv = all([module.can_cache_kv for module in self.modules() if isinstance(module, Attention)])
+
     def forward(
         self,
         x,
@@ -118,6 +123,7 @@ class ContinuousTransformerWrapper(Module):
         return_intermediates = False,
         return_mems = False,
         mask = None,
+        lens = None,
         return_attn = False,
         mems = None,
         mem_masks = None,
@@ -128,6 +134,16 @@ class ContinuousTransformerWrapper(Module):
     ):
         batch, seq, orig_mask, device = *x.shape[:2], mask, x.device
 
+        # maybe seq lengths passed in
+
+        if exists(lens):
+            assert not exists(mask), 'either `mask` or `lens` passed in, but not both'
+            seq_arange = torch.arange(seq, device = device)
+
+            mask = einx.less('j, i -> i j', seq_arange, lens)
+
+        # project in + positional embedding
+
         x = self.project_in(x)
         x = x + self.pos_emb(x, pos = pos)
 
@@ -180,7 +196,7 @@ class ContinuousTransformerWrapper(Module):
         if not return_embeddings and self.probabilistic:
             mean, log_var = rearrange(out, '... (d mean_log_var) -> mean_log_var ... d', mean_log_var = 2)
             variance = log_var.exp()
-            return stack((mean, variance))
+            out = stack((mean, variance))
 
         if return_intermediates:
             return out, intermediates
@@ -223,9 +239,12 @@ class ContinuousAutoregressiveWrapper(Module):
         start_tokens,
         seq_len,
         temperature = 1.,
+        cache_kv = True,
         **kwargs
     ):
+        should_cache_kv = cache_kv and self.net.can_cache_kv
         device = start_tokens.device
+
         was_training = self.net.training
         num_dims = len(start_tokens.shape)
 
@@ -239,10 +258,14 @@ class ContinuousAutoregressiveWrapper(Module):
         self.net.eval()
         out = start_tokens
 
+        cache = None
+
         for _ in range(seq_len):
             x = out[:, -self.max_seq_len:]
 
-            last_output = self.net(x, **kwargs)[..., -1:, :]
+            net_out, new_cache = self.net(x, cache = cache, return_intermediates = True, **kwargs)
+
+            last_output = net_out[..., -1:, :]
 
             if self.probabilistic:
                 mean, var = last_output
@@ -250,6 +273,9 @@ class ContinuousAutoregressiveWrapper(Module):
 
             out = cat((out, last_output), dim = -2)
 
+            if should_cache_kv:
+                cache = new_cache
+
         out = out[:, t:]
 
         if num_dims == 2:
@@ -268,6 +294,7 @@ class ContinuousAutoregressiveWrapper(Module):
         assert 'prepend_embeds' not in kwargs
 
         mask = kwargs.get('mask', None)
+
         if exists(mask) and mask.shape[1] == x.shape[1]:
             mask = mask[:, :-1]
             kwargs['mask'] = mask