x-transformers 2.3.15__py3-none-any.whl → 2.3.16__py3-none-any.whl

x_transformers/continuous.py

@@ -130,6 +130,7 @@ class ContinuousTransformerWrapper(Module):
         sum_embeds = None,
         prepend_embeds = None,
         prepend_mask = None,
+        seq_start_pos = None,
         **kwargs
     ):
         batch, seq, orig_mask, device = *x.shape[:2], mask, x.device
@@ -145,7 +146,7 @@ class ContinuousTransformerWrapper(Module):
         # project in + positional embedding
 
         x = self.project_in(x)
-        x = x + self.pos_emb(x, pos = pos)
+        x = x + self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos)
 
         if exists(sum_embeds):
             x = x + sum_embeds
@@ -315,6 +316,8 @@ class ContinuousAutoregressiveWrapper(Module):
 
         # get target
 
+        seq_start_pos = None
+
         if one_step_autoregress:
             target = x[:, None, 1:]
         else:
@@ -328,6 +331,8 @@ class ContinuousAutoregressiveWrapper(Module):
 
             target = x[batch_arange, target_indices] # rollout targets
 
+            seq_start_pos = torch.zeros(batch, device = device, dtype = torch.long)
+
         # assert inputs
 
         assert 'prepend_embeds' not in kwargs
@@ -363,12 +368,15 @@ class ContinuousAutoregressiveWrapper(Module):
 
             # forward
 
-            out = self.net(inp, mask = step_mask, **kwargs)
+            out = self.net(inp, mask = step_mask, seq_start_pos = seq_start_pos, **kwargs)
 
             outputs.append(out)
 
             inp = out
 
+            if not one_step_autoregress:
+                seq_start_pos.sub_(1)
+
         # stack masks and predictions from rollouts
 
         masks = stack(masks, dim = 1) if exists(mask) else None

{x_transformers-2.3.15.dist-info → x_transformers-2.3.16.dist-info}/METADATA

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

{x_transformers-2.3.15.dist-info → x_transformers-2.3.16.dist-info}/RECORD

@@ -2,7 +2,7 @@ x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,98
 x_transformers/attend.py,sha256=xFsBtl7h7_qebPh7kE81BpmCWAjCgFpB9i_IHu_91es,17288
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
-x_transformers/continuous.py,sha256=5SZmi3Bd77aJAu50f4y1OwwruZd_3ZHptC8dtQmvvxM,11387
+x_transformers/continuous.py,sha256=jy2wsQ3sS80Qwm_gnAmdAnzBfzLoWrGPacOTzU1Q6JM,11674
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
@@ -11,7 +11,7 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 x_transformers/x_transformers.py,sha256=ZfOXrZSiy2jlZ8wVmDdMTLW4hAY_qfmPQHW9t2ABxbo,114097
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.15.dist-info/METADATA,sha256=Dv6r9tZhbF-_7q5yhgmo3pto-D6NVJnYPNnSBEBt73I,89897
-x_transformers-2.3.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.15.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.15.dist-info/RECORD,,
+x_transformers-2.3.16.dist-info/METADATA,sha256=-lL73g4mG5pszuaU7lPdMVGJ7ZtqBqhaejr5VvWWUiw,89897
+x_transformers-2.3.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.16.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.16.dist-info/RECORD,,