x-transformers 1.37.3__py3-none-any.whl → 1.37.4__py3-none-any.whl

x_transformers/nonautoregressive_wrapper.py

@@ -309,9 +309,11 @@ class NonAutoregressiveWrapper(nn.Module):
         with context():
             logits = self.net(masked, **kwargs)
 
+        loss_fn = F.cross_entropy if not self.net.is_log_prob else F.nll_loss
+
         # cross entropy loss
 
-        loss = F.cross_entropy(
+        loss = loss_fn(
             logits[mask],
             orig_seq[mask]
         )

x_transformers/xl_autoregressive_wrapper.py

@@ -40,7 +40,7 @@ class XLAutoregressiveWrapper(nn.Module):
         eos_token = None,
         temperature = 1.,
         filter_logits_fn = top_k,
-        filter_thres = 0.9,
+        filter_kwargs: dict = dict(),
         mems = None,
         **kwargs
     ):
@@ -88,7 +88,7 @@ class XLAutoregressiveWrapper(nn.Module):
             mems = cache.mems
 
             logits = logits[:, -1]
-            filtered_logits = filter_logits_fn(logits, thres = filter_thres)
+            filtered_logits = filter_logits_fn(logits, **filter_kwargs)
             probs = F.softmax(filtered_logits / temperature, dim=-1)
 
             sample = torch.multinomial(probs, 1)
@@ -131,7 +131,9 @@ class XLAutoregressiveWrapper(nn.Module):
 
         split_x = x.split(max_seq_len, dim = -1)
         split_labels = labels.split(max_seq_len, dim = -1)
-        loss_weights = tuple(map(lambda t: t.shape[-1] / seq_len, split_x))
+        loss_weights = tuple((t.shape[-1] / seq_len) for t in split_x)
+
+        loss_fn = F.cross_entropy if not self.net.is_log_prob else F.nll_loss
 
         # go through each chunk and derive weighted losses
 
@@ -146,7 +148,7 @@ class XLAutoregressiveWrapper(nn.Module):
                 **kwargs
             )
 
-            loss = F.cross_entropy(
+            loss = loss_fn(
                 rearrange(logits, 'b n c -> b c n'),
                 chunk_labels,
                 ignore_index = ignore_index

{x_transformers-1.37.3.dist-info → x_transformers-1.37.4.dist-info}/METADATA

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

{x_transformers-1.37.3.dist-info → x_transformers-1.37.4.dist-info}/RECORD

@@ -4,12 +4,12 @@ x_transformers/autoregressive_wrapper.py,sha256=2FN4ZobFcdDGDGWEnUof_geb16dRGSJy
 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=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
+x_transformers/nonautoregressive_wrapper.py,sha256=FDSNVq7wkq4c9muHKngaVPvpHfISfqA8MI7fZZ0KcOY,10503
 x_transformers/x_transformers.py,sha256=gOJBZzOJMu5RkIsxw9TZtde4Sx--D18yX8LjrYIsPbE,83677
-x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
+x_transformers/xl_autoregressive_wrapper.py,sha256=mwygaOG6ECJGrzG1ViA6Xxu9m_WSYOzPUDzUhkt4m0U,4188
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.37.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.37.3.dist-info/METADATA,sha256=SIGTCQMrLkyq_aksJAst0iXw9VfFT6QWlGvtUElbTMg,661
-x_transformers-1.37.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.37.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.37.3.dist-info/RECORD,,
+x_transformers-1.37.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.37.4.dist-info/METADATA,sha256=H1oc8hkvPpEoS4ioSOn1wVSHjcTJ3OJcFdJuJuuWdl4,661
+x_transformers-1.37.4.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+x_transformers-1.37.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.37.4.dist-info/RECORD,,