x-transformers 2.3.1__py3-none-any.whl → 2.3.2__py3-none-any.whl

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,
@@ -180,7 +185,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 +228,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 +247,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 +262,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:

{x_transformers-2.3.1.dist-info → x_transformers-2.3.2.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,98
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
-x_transformers/continuous.py,sha256=F5XPQU5Y798R1_JoepX4Mg44_j3Whw8SHaTsavq1YZs,8256
+x_transformers/continuous.py,sha256=yAE8hLyusrEd-12mkgLASDL-cFgpZQf32s93FKfez7o,8674
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
@@ -11,7 +11,7 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 x_transformers/x_transformers.py,sha256=MF91aJGr2DOjIGe57uqwgyNxCExBg_tI9z7usAJMxOM,112401
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.3.1.dist-info/METADATA,sha256=-y3iEikqisIdIx8eBfP41qVZj2Nqzpm88usIUek6Pwg,88686
-x_transformers-2.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.1.dist-info/RECORD,,
+x_transformers-2.3.2.dist-info/METADATA,sha256=8m6hpJlMKesI-SLxth_9z0VYIHUh7bTWsJ9Am5OSni4,88686
+x_transformers-2.3.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.2.dist-info/RECORD,,