x-transformers 1.24.6__py3-none-any.whl → 1.25.0__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/continuous.py

@@ -142,11 +142,13 @@ class ContinuousAutoregressiveWrapper(nn.Module):
         self,
         net: ContinuousTransformerWrapper,
         ignore_index = -100,
-        pad_value = 0
+        pad_value = 0,
+        loss_fn = nn.MSELoss(reduction = 'none')
     ):
         super().__init__()
         self.net = net
         self.max_seq_len = net.max_seq_len
+        self.loss_fn = loss_fn
 
     @torch.no_grad()
     def generate(self, start_tokens, seq_len, **kwargs):
@@ -187,9 +189,11 @@ class ContinuousAutoregressiveWrapper(nn.Module):
             kwargs['mask'] = mask
 
         out = self.net(inp, **kwargs)
-        loss = F.mse_loss(out, target, reduction = 'none')
+
+        loss = self.loss_fn(out, target)
 
         if exists(mask):
+            assert loss.ndim > 1, 'loss should not be reduced if mask is passed in'
             loss = loss[mask]
 
         return loss.mean()

x_transformers/x_transformers.py

@@ -1331,6 +1331,33 @@ 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_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_len):
+            prefix_mask = torch.arange(n, device = device) < rearrange(prefix_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.6.dist-info → x_transformers-1.25.0.dist-info}/METADATA

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

x_transformers-1.25.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+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=btyWp8_gOX8jkTqVUWQvuVaE8x6R8IMEkM6_Nczxsc0,59402
+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.24.6.dist-info/RECORD

@@ -1,13 +0,0 @@
-x_transformers/__init__.py,sha256=xeZcV-MAaT-l8NOktPa_54WkmXLn4nG5E7CyKQC2iAI,610
-x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,10188
-x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
-x_transformers/continuous.py,sha256=mLVaSaPSdZyB6ctIP54ax5SMmUo3vTWhp5KgXJpKuns,5349
-x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=exkQkU_3S5RL0OA9tX4-h8ZWPOdGkoszvLkOStYVj_o,58543
-x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers/xval.py,sha256=lS9W_E_RskPQAqVZkPiUzbByoW1Ajsw_phsikA3JDAg,8139
-x_transformers-1.24.6.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.24.6.dist-info/METADATA,sha256=9HcuKlXrBntL3TZIb1vyDRqj8GleTyENK5hiFJrZ7jg,661
-x_transformers-1.24.6.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-x_transformers-1.24.6.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.24.6.dist-info/RECORD,,