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

x_transformers/__init__.py

@@ -2,6 +2,7 @@ from x_transformers.x_transformers import (
     XTransformer,
     Encoder,
     Decoder,
+    PrefixDecoder,
     CrossAttender,
     Attention,
     TransformerWrapper,

x_transformers/x_transformers.py

@@ -1331,6 +1331,36 @@ class Decoder(AttentionLayers):
         assert 'causal' not in kwargs, 'cannot set causality on decoder'
         super().__init__(causal = True, **kwargs)
 
+class PrefixDecoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert 'causal' not in kwargs, 'cannot set causality on decoder'
+        super().__init__(causal = False, **kwargs)
+
+    def forward(
+        self,
+        x,
+        *args,
+        attn_mask = None,
+        prefix_attn_len = None,
+        **kwargs
+    ):
+        b, n, device = *x.shape[:2], x.device
+        causal_mask = torch.ones((n, n), device = device, dtype = torch.bool).triu(1)
+
+        forwarded_mask = ~causal_mask
+
+        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):
+            forwarded_mask = forwarded_mask & attn_mask
+
+        return super().forward(x, *args, attn_mask = forwarded_mask, **kwargs)
+
 class CrossAttender(AttentionLayers):
     def __init__(self, **kwargs):
         super().__init__(cross_attend = True, only_cross = True, **kwargs)

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

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

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

@@ -1,13 +1,13 @@
-x_transformers/__init__.py,sha256=xeZcV-MAaT-l8NOktPa_54WkmXLn4nG5E7CyKQC2iAI,610
+x_transformers/__init__.py,sha256=pXc_U4M3ONUQcpNgZySDIlCF1rp7u4FFmcOYjc4WuXw,629
 x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,10188
 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=exkQkU_3S5RL0OA9tX4-h8ZWPOdGkoszvLkOStYVj_o,58543
+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.24.7.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.24.7.dist-info/METADATA,sha256=qLvJYXptB_smV5-qKfv1vU-bhuz9Uh381fKrKfOA6i0,661
-x_transformers-1.24.7.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-x_transformers-1.24.7.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.24.7.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,,