x-transformers 2.1.16__py3-none-any.whl → 2.1.17__py3-none-any.whl

x_transformers/belief_state_wrapper.py

@@ -23,7 +23,7 @@ from x_transformers.x_transformers import (
 )
 
 import einx
-from einops import rearrange, repeat
+from einops import rearrange, repeat, pack, unpack
 from einops.layers.torch import Rearrange
 
 # helper functions
@@ -69,6 +69,8 @@ class BeliefStateWrapper(Module):
         dim = forward_decoder.emb_dim
         num_tokens = forward_decoder.num_tokens
 
+        self.to_forward_logits = nn.Linear(dim, num_tokens, bias = False)
+
         # the suffix token
 
         self.suffix_token = nn.Parameter(torch.zeros(dim))
@@ -126,7 +128,7 @@ class BeliefStateWrapper(Module):
         prompts,
         seq_len,
         temperature = 1.25,
-        cache_kv = True,
+        cache_kv = False,
         suffix: Tensor | None = None, # the goal conditioning
         filter_logits_fn = min_p,
         filter_kwargs = dict(
@@ -136,6 +138,8 @@ class BeliefStateWrapper(Module):
     ):
         max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
 
+        prompts, batch_ps = pack([prompts], '* d')
+
         batch, orig_seq_len = prompts.shape
 
         out = prompts
@@ -197,14 +201,19 @@ class BeliefStateWrapper(Module):
 
             # concat sample
 
-            out = torch.cat((out, sample), dim=-1)
+            out = torch.cat((out, sample), dim = -1)
+
+        out = out[:, orig_seq_len:]
 
-        return out[:, orig_seq_len:]
+        out, = unpack(out, batch_ps, '* n')
+
+        return out
 
     def forward(
         self,
         seq,
-        backward = True
+        return_loss_only = False,
+        loss_scale = 1.
     ):
         batch, seq_len, device = *seq.shape, seq.device
 
@@ -244,6 +253,7 @@ class BeliefStateWrapper(Module):
         # f - forward, b - backward, i - indices
 
         fi, bi = fb_pairs.unbind(dim = -1)
+
         valid_mask = (bi - fi) >= 2
 
         fb_pairs = fb_pairs[valid_mask]
@@ -265,8 +275,9 @@ class BeliefStateWrapper(Module):
 
         labels_fi, labels_bi = (fi + 1), bi
 
-        forward_labels, backward_labels = seq[:, fi], seq[:, bi]
-        labels = stack((forward_labels, backward_labels), dim = -1)
+        forward_labels, backward_labels = seq[:, labels_fi], seq[:, labels_bi]
+
+        labels = cat((forward_labels, backward_labels), dim = -1)
 
         # get the forward and backward embedding pairs and feed them through the text head for both forward and backward predictions
 
@@ -281,7 +292,7 @@ class BeliefStateWrapper(Module):
 
         loss = F.cross_entropy(
             rearrange(logits, 'b n (fb l) -> b l (fb n)', fb = 2),
-            rearrange(labels, 'b n fb -> b (fb n)'),
+            labels,
             reduction = 'none' if self.needs_loss_weight else 'mean'
         )
 
@@ -290,12 +301,12 @@ class BeliefStateWrapper(Module):
         if exists(self.to_terminal_logit):
             terminal_logits = self.to_terminal_logit(fb_embeds)
 
-            labels = ((bi - fi) == 2).float() # distance is exactly 2
-            labels = repeat(labels, 'n -> b n', b = batch)
+            terminal_labels = ((bi - fi) == 2).float() # distance is exactly 2
+            terminal_labels = repeat(terminal_labels, 'n -> b n', b = batch)
 
             is_end_loss = F.binary_cross_entropy_with_logits(
                 terminal_logits,
-                labels
+                terminal_labels
             )
 
             loss = (
@@ -303,6 +314,11 @@ class BeliefStateWrapper(Module):
                 is_end_loss * self.pred_terminal_loss_weight
             )
 
+        # maybe early return loss
+
+        if return_loss_only:
+            return loss
+
         # maybe loss weighting
 
         if self.needs_loss_weight:
@@ -312,18 +328,9 @@ class BeliefStateWrapper(Module):
 
         # backwards
 
-        orig_backward = getattr(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
+        (loss * loss_scale).backward()
 
-        if backward:
-            patched_backward_fn()
-        else:
-            setattr(loss, 'backward', patched_backward_fn)
+        orig_forward_embeds.backward(forward_embeds.grad)
+        orig_backward_embeds.backward(backward_embeds.grad)
 
         return loss

{x_transformers-2.1.16.dist-info → x_transformers-2.1.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.16
+Version: 2.1.17
 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.16.dist-info → x_transformers-2.1.17.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=VjUB73yFBWevN6xMc6_1s-Yc58pJv8SDAUUEXwpR-W0,9842
+x_transformers/belief_state_wrapper.py,sha256=7zi_inz75jZVQE4zFVA2H45yR2V0RLTgX_aGO86rN9s,9878
 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.16.dist-info/METADATA,sha256=MXaag1fuq1BsAmyh9k8sSRiHpy-jAUQW2Hn1GC53MnQ,87571
-x_transformers-2.1.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.16.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.16.dist-info/RECORD,,
+x_transformers-2.1.17.dist-info/METADATA,sha256=PzCTyZz0yqv6vXVpzZwtWyu6sjR-kKwcgVPAsrn5TnI,87571
+x_transformers-2.1.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.17.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.17.dist-info/RECORD,,