x-transformers 2.1.22__py3-none-any.whl → 2.1.24__py3-none-any.whl

x_transformers/belief_state_wrapper.py

@@ -228,7 +228,8 @@ class BeliefStateWrapper(Module):
         self,
         seq,
         return_loss_only = False,
-        loss_scale = 1.
+        loss_scale = 1.,
+        loss_weight_by_fb_indices: callable | None = None
     ):
         batch, seq_len, device = *seq.shape, seq.device
 
@@ -336,9 +337,27 @@ class BeliefStateWrapper(Module):
 
         # maybe loss weighting
 
-        if self.needs_loss_weight:
+        needs_loss_weight = default(self.needs_loss_weight, exists(loss_weight_by_fb_indices))
+
+        if needs_loss_weight:
             loss = rearrange(loss, 'b (fb n) -> b fb n', fb = 2)
-            loss = einx.multiply('b fb n, fb', loss, self.loss_weights)
+
+            if self.needs_loss_weight:
+                loss = einx.multiply('b fb n, fb', loss, self.loss_weights)
+
+            # allow researcher to pass in a function that acts on the the forward backward indices Int['n fb']
+            # the reason this may be needed is because the earlier tokens will have more eligible pairs for training, and perhaps this could be normalized
+
+            if exists(loss_weight_by_fb_indices):
+                loss_weight = loss_weight_by_fb_indices(fb_pairs)
+
+                if loss_weight.ndim == 1:
+                    loss = einx.multiply('b fb n, n', loss, loss_weight)
+                elif loss_weight.ndim == 2:
+                    loss = einx.multiply('b fb n, n fb', loss, loss_weight)
+                else:
+                    raise ValueError('invalid loss weight dims')
+
             loss = loss.mean()
 
         # backwards

{x_transformers-2.1.22.dist-info → x_transformers-2.1.24.dist-info}/METADATA

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

{x_transformers-2.1.22.dist-info → x_transformers-2.1.24.dist-info}/RECORD

@@ -1,7 +1,7 @@
 x_transformers/__init__.py,sha256=NDoiBivau559WQ0FvXG4ssU3Il9aoHmTIUFN_1juz0s,911
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
-x_transformers/belief_state_wrapper.py,sha256=crj0yaTNmszDYlueGu_plGKpxVg0GKH8Z-B66SQluNs,10550
+x_transformers/belief_state_wrapper.py,sha256=S_XP0RQpXemGibYLrVHgNAm0I0IqeocgSMlShcwRqG8,11452
 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
@@ -10,7 +10,7 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 x_transformers/x_transformers.py,sha256=fqgtIs6__JpLWMnJa8AY5OW3AJ2GR1B5p-9TsWdiOIU,110425
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.1.22.dist-info/METADATA,sha256=FBxTg2dObipuXg2cC_ykYpzLF1AHEebrMyRjoAc0Xk4,87875
-x_transformers-2.1.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.22.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.22.dist-info/RECORD,,
+x_transformers-2.1.24.dist-info/METADATA,sha256=sl2x1DgYJP3FMPO4cZiSJYCISkgrNA1flgdeVYFl3GA,87875
+x_transformers-2.1.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.24.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.24.dist-info/RECORD,,