x-transformers 2.3.24__py3-none-any.whl → 2.3.26__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -309,7 +309,13 @@ class AutoregressiveWrapper(Module):
 
         return out
 
-    def forward(self, x, return_outputs = False, **kwargs):
+    def forward(
+        self,
+        x,
+        return_outputs = False,
+        prepend_embeds = None,
+        **kwargs
+    ):
         seq, ignore_index, add_attn_z_loss, add_next_embed_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss, self.add_continuous_pred_head
 
         inp, target = x, x[:, 1:]
@@ -328,6 +334,7 @@ class AutoregressiveWrapper(Module):
             return_intermediates = True,
             return_attn_z_loss = add_attn_z_loss,
             return_next_embed_pred = add_next_embed_loss,
+            prepend_embeds = prepend_embeds,
             **kwargs
         )
 
@@ -338,6 +345,14 @@ class AutoregressiveWrapper(Module):
         else:
             logits = out
 
+        # if there are prepended embeds, excise it out
+
+        if exists(prepend_embeds):
+            prepend_len = prepend_embeds.shape[1]
+            logits = logits[:, prepend_len:]
+
+        # take all tokens but the last
+
         logits = logits[:, :-1]
 
         # loss function
@@ -356,7 +371,7 @@ class AutoregressiveWrapper(Module):
             loss = loss + cache.attn_z_loss
 
         if add_next_embed_loss:
-            mask = inp[:, :-1] != ignore_index
+            mask = target != ignore_index
             embed_pred = next_embed_pred[:, :-1]
             cont_targets = init_embeds[:, 1:].detach()
 

{x_transformers-2.3.24.dist-info → x_transformers-2.3.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.3.24
+Version: 2.3.26
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.3.24.dist-info → x_transformers-2.3.26.dist-info}/RECORD

@@ -1,6 +1,6 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
 x_transformers/attend.py,sha256=Ax34Rw56xXAWfFPqtZ_B8iKEW2EfQdbVoc9uFjfeNjA,17404
-x_transformers/autoregressive_wrapper.py,sha256=tMVbIC8iXTpfGDxRhPqqHTulkxB8aZqNML77WbGhfac,11755
+x_transformers/autoregressive_wrapper.py,sha256=3tDUiY5kNcxNUjRERoeuFV0mXztOvgGrckoACIfHvqI,12091
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=hpb1sSbt3k2LNzzjrjSd8F5xOIbKj7IluV9MBEAFLkw,13031
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
@@ -11,7 +11,7 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 x_transformers/x_transformers.py,sha256=7phSZvP1_SDRIkVMwVR4cz1dFU2UlR2Wf1HJHEQlcQg,116222
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.24.dist-info/METADATA,sha256=vqW6_PFF3JiQirofvdzEMXwAt_x9luG_TOQimtIWwg8,89897
-x_transformers-2.3.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.24.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.24.dist-info/RECORD,,
+x_transformers-2.3.26.dist-info/METADATA,sha256=Qc0zIph59FLOC0GPGIe41M6P1SD_lljzKg5ytoMyPAI,89897
+x_transformers-2.3.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.26.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.26.dist-info/RECORD,,