x-transformers 1.38.2__py3-none-any.whl → 1.38.3__py3-none-any.whl

x_transformers/attend.py

@@ -67,7 +67,7 @@ print_once = once(print)
 
 # selective attention
 # https://arxiv.org/abs/2410.02703 - section 3.3
-# it is a parameter-less technique to allow each token to prevent itself from being attended to by future tokens
+# it is a technique to allow each token to prevent itself from being attended to by future tokens
 # if sim_head_gate not supplied, will use the first head of the attention logits (sim in this framework)
 
 def selective_attn(
@@ -81,9 +81,10 @@ def selective_attn(
     gate = F.relu(sim_head_gate) # only positive
 
     if no_mask_sos:
+        gate = gate.clone()
         gate[..., -i] = 0.
 
-    eye = torch.eye(j, device = device)
+    eye = torch.eye(i, device = device)
 
     if j > i:
         eye = F.pad(eye, (j - i, 0), value = 1.)

x_transformers/x_transformers.py

@@ -1084,6 +1084,10 @@ class Attention(Module):
 
         self.rotary_embed_values = rotary_embed_values
 
+        # whether parent can kv cache
+
+        self.can_cache_kv = not selective
+
         # init output projection 0
 
         if zero_init_output:
@@ -1634,6 +1638,10 @@ class AttentionLayers(Module):
                 residual
             ]))
 
+        # determine whether 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,
@@ -2135,7 +2143,7 @@ class TransformerWrapper(Module):
 
         # whether can do cached kv decoding
 
-        self.can_cache_kv = self.num_memory_tokens == 0 and not recycling
+        self.can_cache_kv = self.num_memory_tokens == 0 and not recycling and self.attn_layers.can_cache_kv
         self.can_cache_kv_outside_max_seq_len = no_abs_pos_emb
 
     def init_(self):

{x_transformers-1.38.2.dist-info → x_transformers-1.38.3.dist-info}/METADATA

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

{x_transformers-1.38.2.dist-info → x_transformers-1.38.3.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=HO6HZ1fowJ6a6v915PH4s8PnfNj0_q47Sq7yc9AP5YQ,15380
+x_transformers/attend.py,sha256=Z9cnbY3f9zCl1yUOMWVQT5_ee4keC0lo_NX4cd0rbKk,15393
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=Dol6GMZOoHGOFdHwe21o2SbJp6b3YKCUHoIs_AjfvTo,83963
+x_transformers/x_transformers.py,sha256=P3o1R2DY2ic71FlJJ4ie4w_z-g3jIIrkBXcbllRoXHA,84240
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.38.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.38.2.dist-info/METADATA,sha256=pThqFTEo8bihgUlSYdv3r-JFB153pbO2baJgXwMMZZs,661
-x_transformers-1.38.2.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.38.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.38.2.dist-info/RECORD,,
+x_transformers-1.38.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.38.3.dist-info/METADATA,sha256=XbmJ91NWqmOjO0xvyXh9uh82TIwVbU54L3eFmaFVvYs,661
+x_transformers-1.38.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+x_transformers-1.38.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.38.3.dist-info/RECORD,,