x-transformers 1.30.3__py3-none-any.whl → 1.30.6__py3-none-any.whl

x_transformers/attend.py

@@ -22,6 +22,7 @@ class Intermediates:
     pre_softmax_attn:   Tensor | None = None
     post_softmax_attn:  Tensor | None = None
     cached_kv:          Tuple[Tensor, Tensor] | None = None
+    layer_type:         str | None = None
 
     def to_tuple(self):
         return (self.qk_similarities, self.pre_softmax_attn, self.post_softmax_attn)

x_transformers/x_transformers.py

@@ -868,7 +868,7 @@ class Attention(Module):
 
         k, v, r = map(lambda t: maybe(rearrange)(t, 'b n (h d) -> b h n d', h = kv_h), (k, v, r))
 
-        if exists(cache) and not has_context:
+        if exists(cache):
             ck, cv = cache.cached_kv
 
             if exists(mem):
@@ -1338,6 +1338,9 @@ class AttentionLayers(Module):
         if exists(cache):
             assert not self.training and self.causal and not any([*map(exists, (mask, attn_mask))])
 
+            if exists(context):
+                context = context[:, :0]
+
             if cache_age > 0:
                 x = x[:, -cache_age:] # for spec decoding, may be greater than 1
 
@@ -1407,6 +1410,7 @@ class AttentionLayers(Module):
             x = residual_fn(out, inner_residual)
 
             if layer_type in ('a', 'c') and return_hiddens:
+                inter.layer_type = layer_type
                 intermediates.append(inter)
 
             if layer_type == 'a' and self.residual_attn:

{x_transformers-1.30.3.dist-info → x_transformers-1.30.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.30.3
+Version: 1.30.6
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.30.3.dist-info → x_transformers-1.30.6.dist-info}/RECORD

@@ -1,14 +1,14 @@
 x_transformers/__init__.py,sha256=8LQl-dNL6vj8VHRx5LMSOlRDTXQvYOuM21PDXz8WdiI,703
-x_transformers/attend.py,sha256=2SPHjXS_QAAZt04lHWGtdOypTExmo3BrbFhgcIQTk-Y,10671
+x_transformers/attend.py,sha256=8opOeCQddi440WcH73B_wB5vtL0jaEQwBL-DIWq2lCs,10713
 x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3nPqO6xp2FFCc,9619
 x_transformers/continuous.py,sha256=WO52n9lFAXv5-SGadi2cApGF8dkouN8QSTEOuC7erj8,6180
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=BQypGJAoqXrAe_ek95wUcXdSQdAWjvw5mEti-H1JxcI,66288
+x_transformers/x_transformers.py,sha256=pXckFcDL6kTghYEUjIamZiR5H8dV6aIEPQTIYAGgqxA,66388
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.30.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.30.3.dist-info/METADATA,sha256=FdVMtNhhggibbG3fJfCfGdRUAp6fLkTyG8KzDUq_r1Y,661
-x_transformers-1.30.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-x_transformers-1.30.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.30.3.dist-info/RECORD,,
+x_transformers-1.30.6.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.30.6.dist-info/METADATA,sha256=1Xq9oSctaCQ5TOpdM3j6lJENYStLuda5VqEzOtq1B0c,661
+x_transformers-1.30.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+x_transformers-1.30.6.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.30.6.dist-info/RECORD,,