x-transformers 2.1.5__py3-none-any.whl → 2.1.7__py3-none-any.whl

x_transformers/__init__.py

@@ -14,6 +14,7 @@ from x_transformers.x_transformers import (
 
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
 from x_transformers.nonautoregressive_wrapper import NonAutoregressiveWrapper
+from x_transformers.belief_state_wrapper import BeliefStateWrapper
 
 from x_transformers.continuous import (
     ContinuousTransformerWrapper,

x_transformers/{belief_state.py → belief_state_wrapper.py}

@@ -6,7 +6,7 @@
 import torch
 from torch.autograd import Function
 from torch.nn import Module, ModuleList
-from torch import nn, cat, stack, arange, cartesian_prod
+from torch import nn, cat, stack, tensor, arange, cartesian_prod
 import torch.nn.functional as F
 
 from x_transformers.x_transformers import (
@@ -36,7 +36,8 @@ class BeliefStateWrapper(Module):
         self,
         forward_decoder: TransformerWrapper,
         backward_decoder: TransformerWrapper,
-        train_frac_forward_backward_pairs: float = 1.
+        train_frac_forward_backward_pairs: float = 1.,
+        backward_ar_loss_weight: float = 1. # can weigh the training of the backwards decoder differently, perhaps fwd/bwd have a shared backbone etc etc
     ):
         super().__init__()
         assert forward_decoder.emb_dim == backward_decoder.emb_dim, 'forward and backwards model must have the same embedding dimension'
@@ -70,9 +71,17 @@ class BeliefStateWrapper(Module):
         self.train_frac_fb_pairs = train_frac_forward_backward_pairs
         self.needs_subsample_fb_pairs = train_frac_forward_backward_pairs < 1.
 
+        # loss weighting
+
+        self.backward_ar_loss_weight = backward_ar_loss_weight
+        self.needs_loss_weight = backward_ar_loss_weight != 1.
+
+        self.register_buffer('loss_weights', tensor([1., self.backward_ar_loss_weight]))
+
     def forward(
         self,
-        seq
+        seq,
+        backward = True
     ):
         batch, seq_len, device = *seq.shape, seq.device
 
@@ -149,14 +158,31 @@ class BeliefStateWrapper(Module):
 
         fb_loss = F.cross_entropy(
             rearrange(logits, 'b n (fb l) -> b l (fb n)', fb = 2),
-            rearrange(labels, 'b n fb -> b (fb n)')
+            rearrange(labels, 'b n fb -> b (fb n)'),
+            reduction = 'none' if self.needs_loss_weight else 'mean'
         )
 
+        # maybe loss weighting
+
+        if self.needs_loss_weight:
+            fb_loss = rearrange(fb_loss, 'b (fb n) -> b fb n', fb = 2)
+            fb_loss = einx.multiply('b fb n, fb', fb_loss, self.loss_weights)
+            fb_loss = fb_loss.mean()
+
         # backwards
 
-        fb_loss.backward()
+        orig_backward = getattr(fb_loss, 'backward')
+
+        def patched_backward_fn(*args, **kwargs):
+            orig_backward(*args, **kwargs)
+            orig_forward_embeds.backward(forward_embeds.grad)
+            orig_backward_embeds.backward(backward_embeds.grad)
+
+        # can allow the researcher to call .backward from the outside
 
-        orig_forward_embeds.backward(forward_embeds.grad)
-        orig_backward_embeds.backward(backward_embeds.grad)
+        if backward:
+            patched_backward_fn()
+        else:
+            setattr(fb_loss, 'backward', patched_backward_fn)
 
         return fb_loss

{x_transformers-2.1.5.dist-info → x_transformers-2.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.5
+Version: 2.1.7
 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.5.dist-info → x_transformers-2.1.7.dist-info}/RECORD

@@ -1,7 +1,7 @@
-x_transformers/__init__.py,sha256=l0dom8ZYkRzFvnDdgzDboXqrI1tKav3beVE7TN2nHko,844
+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.py,sha256=GfYDeDqmhldozECgFsJ9zhd6O5NMvdYA5OwueVs8SB4,4742
+x_transformers/belief_state_wrapper.py,sha256=WoagJe_cHguWAWzhEAmeDD_TzUsQcMQkgP2Mt88wrAg,5827
 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=-80N4sqUr3sR51Ms4wCfc4jhxnPwf0ApNR4xfIsasfQ,110142
 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.5.dist-info/METADATA,sha256=-nQpm1eBGBXkLEXGiQK06NaIWa13CEnBujBNfTzvnJ8,87570
-x_transformers-2.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.5.dist-info/RECORD,,
+x_transformers-2.1.7.dist-info/METADATA,sha256=D7enhdRYdhTddt4kGSS-rSsvtHgtFaLYQH-cVuYxUVM,87570
+x_transformers-2.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.7.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.7.dist-info/RECORD,,