x-transformers 1.42.10__py3-none-any.whl → 1.42.12__py3-none-any.whl

x_transformers/continuous.py

@@ -2,7 +2,8 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
-from einops import pack, repeat, unpack
+import einx
+from einops import reduce, pack, repeat, unpack
 
 from x_transformers.x_transformers import (
     AttentionLayers,
@@ -24,6 +25,15 @@ def default(val, d):
         return val
     return d() if callable(d) else d
 
+def masked_mean(t, mask):
+    t = einx.where('b n, b n d, -> b n d', mask, t, 0.)
+
+    num = reduce(t, 'b n d -> b', 'sum')
+    den = mask.sum(dim = -1)
+
+    masked_average = num / den.clamp(min = 1.)
+    return masked_average
+
 # main classes
 
 class ContinuousTransformerWrapper(nn.Module):
@@ -169,12 +179,15 @@ class ContinuousAutoregressiveWrapper(nn.Module):
         net: ContinuousTransformerWrapper,
         ignore_index = -100,
         pad_value = 0,
-        loss_fn = nn.MSELoss(reduction = 'none')
+        loss_fn = nn.MSELoss(reduction = 'none'),
+        equal_loss_weight_batch = False  # setting this to True, if the mask is passed in and sequences are variable in length, each sequence will be weighted the same (as opposed to each token)
     ):
         super().__init__()
         self.net = net
         self.max_seq_len = net.max_seq_len
+
         self.loss_fn = loss_fn
+        self.equal_loss_weight_batch = equal_loss_weight_batch
 
     @torch.no_grad()
     def generate(self, start_tokens, seq_len, **kwargs):
@@ -222,6 +235,10 @@ class ContinuousAutoregressiveWrapper(nn.Module):
 
         if exists(mask):
             assert loss.ndim > 1, 'loss should not be reduced if mask is passed in'
-            loss = loss[mask]
+
+            if self.equal_loss_weight_batch:
+                loss = masked_mean(loss, mask)
+            else:
+                loss = loss[mask]
 
         return loss.mean()

{x_transformers-1.42.10.dist-info → x_transformers-1.42.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.42.10
+Version: 1.42.12
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang
@@ -17,4 +17,5 @@ License-File: LICENSE
 Requires-Dist: torch >=2.0
 Requires-Dist: einx >=0.3.0
 Requires-Dist: einops >=0.8.0
+Requires-Dist: packaging >=21.0
 

{x_transformers-1.42.10.dist-info → x_transformers-1.42.12.dist-info}/RECORD

@@ -1,7 +1,7 @@
 x_transformers/__init__.py,sha256=l0dom8ZYkRzFvnDdgzDboXqrI1tKav3beVE7TN2nHko,844
 x_transformers/attend.py,sha256=SdWlV8Vp5DtpsOzAd0LRhm4VGrJf0lJCGiV2_j_CtoA,17284
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
-x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
+x_transformers/continuous.py,sha256=p0sCAiH1na236ygwgL1Yyhu36eZBf9cZvoW1JyP_fFE,7073
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
@@ -9,8 +9,8 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 x_transformers/x_transformers.py,sha256=VxdA44EYQhVH1Rp7wreJ83I2e0Ea7VN_bFRE-iDXOI8,93833
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.42.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.42.10.dist-info/METADATA,sha256=UlOzLgz1fhOIaz3eKDhcHgVI81EZySt4Ko3BiKJ7Jok,690
-x_transformers-1.42.10.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.42.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.42.10.dist-info/RECORD,,
+x_transformers-1.42.12.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.42.12.dist-info/METADATA,sha256=fC5MiH_nobMaDrmS3CsSAdQEiBXYxJes18Y9GR8L2Ig,722
+x_transformers-1.42.12.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+x_transformers-1.42.12.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.42.12.dist-info/RECORD,,

{x_transformers-1.42.10.dist-info → x_transformers-1.42.12.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.3.0)
+Generator: setuptools (75.5.0)
 Root-Is-Purelib: true
 Tag: py3-none-any