x-transformers 1.25.0__py3-none-any.whl → 1.25.1__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1341,7 +1341,7 @@ class PrefixDecoder(AttentionLayers):
         x,
         *args,
         attn_mask = None,
-        prefix_len = None,
+        prefix_attn_len = None,
         **kwargs
     ):
         b, n, device = *x.shape[:2], x.device
@@ -1349,8 +1349,11 @@ class PrefixDecoder(AttentionLayers):
 
         forwarded_mask = ~causal_mask
 
-        if exists(prefix_len):
-            prefix_mask = torch.arange(n, device = device) < rearrange(prefix_len, 'b -> b 1 1 1')
+        if exists(prefix_attn_len):
+            if isinstance(prefix_attn_len, int):
+                prefix_attn_len = torch.full((b,), prefix_attn_len, device = device)
+
+            prefix_mask = torch.arange(n, device = device) < rearrange(prefix_attn_len, 'b -> b 1 1 1')
             forwarded_mask = forwarded_mask | prefix_mask
 
         if exists(attn_mask):

{x_transformers-1.25.0.dist-info → x_transformers-1.25.1.dist-info}/METADATA

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

{x_transformers-1.25.0.dist-info → x_transformers-1.25.1.dist-info}/RECORD

@@ -3,11 +3,11 @@ x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,1018
 x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
 x_transformers/continuous.py,sha256=G8mVTan2-YbzkY3YDCTar1oyHPMSl0p4F6iRz3Nl0Is,5497
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=btyWp8_gOX8jkTqVUWQvuVaE8x6R8IMEkM6_Nczxsc0,59402
+x_transformers/x_transformers.py,sha256=QF4N6xoVcp8PFMWsRYOMMcqRVjJb5_ArzjPEzWCyffE,59552
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=lS9W_E_RskPQAqVZkPiUzbByoW1Ajsw_phsikA3JDAg,8139
-x_transformers-1.25.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.25.0.dist-info/METADATA,sha256=tIF8EPLHiY2CY0Y0TPFEREgo4VOGDWsv5DLOWi4bnBw,661
-x_transformers-1.25.0.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-x_transformers-1.25.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.25.0.dist-info/RECORD,,
+x_transformers-1.25.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.25.1.dist-info/METADATA,sha256=CnSAhkwS-ZCUYSp0pY1k5PIbIApMw38_75_F1vuof7w,661
+x_transformers-1.25.1.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
+x_transformers-1.25.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.25.1.dist-info/RECORD,,