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

x_transformers/autoregressive_wrapper.py

@@ -36,8 +36,16 @@ def eval_decorator(fn):
 
 # for variable lengthed prefixes
 
+def pad_at_dim(t, pad: tuple[int, int], dim = -1, value = 0.):
+    if pad == (0, 0):
+        return t
+
+    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    zeros = ((0, 0) * dims_from_right)
+    return F.pad(t, (*zeros, *pad), value = value)
+
 def align_right(t, lens, pad_id = 0):
-    batch, seq_len, device, dtype = *t.shape, t.device, t.dtype
+    batch, seq_len, device, dtype = *t.shape[:2], t.device, t.dtype
 
     assert lens.ndim == 1 and lens.shape[0] == batch
     assert lens.amax() <= seq_len
@@ -48,10 +56,10 @@ def align_right(t, lens, pad_id = 0):
     batch_arange = torch.arange(batch, device = device, dtype = torch.long)[..., None]
     prompt_len_arange = torch.arange(seq_len, device = device, dtype = torch.long)
 
-    t = F.pad(t, (max_pad_len, 0), value = pad_id)
+    t = pad_at_dim(t, (max_pad_len, 0), value = pad_id, dim = 1)
     offset = max_pad_len - pad_lens
 
-    aligned = t[batch_arange, prompt_len_arange + offset[..., None]]
+    aligned = t[batch_arange, prompt_len_arange + offset[..., None], ...]
     return aligned
 
 # nucleus

x_transformers/continuous.py

@@ -9,6 +9,8 @@ from torch.distributions import Normal
 import einx
 from einops import rearrange, reduce, pack, repeat, unpack
 
+from x_transformers.autoregressive_wrapper import align_right
+
 from x_transformers.x_transformers import (
     Attention,
     AttentionLayers,
@@ -222,7 +224,6 @@ class ContinuousAutoregressiveWrapper(Module):
         net: ContinuousTransformerWrapper,
         loss_fn: Module | None = 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)
-        rollout_steps = 1                 # they used 2 rollout steps in a successful world model paper https://ai.meta.com/vjepa/
     ):
         super().__init__()
         self.net = net
@@ -236,14 +237,6 @@ class ContinuousAutoregressiveWrapper(Module):
         self.loss_fn = loss_fn
         self.equal_loss_weight_batch = equal_loss_weight_batch
 
-        # num rollout steps - if greater than one, recurrently feedback the output and enforce loss rollout steps - 1 ahead
-        # applied successfully in vjepa2 world model, with rollout steps of 2
-        # rollout steps of 1 would be the same as single step autoregressive
-
-        assert not (rollout_steps > 1 and probabilistic), f'rollout steps greater than 1 only supported for non-probabilistic'
-        assert 1 <= rollout_steps
-        self.rollout_steps = rollout_steps
-
     @torch.no_grad()
     def generate(
         self,
@@ -298,40 +291,17 @@ class ContinuousAutoregressiveWrapper(Module):
         self.net.train(was_training)
         return out
 
-    def forward(
+    def forward_rollout(
         self,
         x,
+        rollout_steps = 2,
         **kwargs
     ):
-        steps = self.rollout_steps
-        one_step_autoregress = steps == 1
-
-        # get the input
-
-        inp = x[:, :-steps]
-
-        # variables
-
-        batch, seq_len, device = *inp.shape[:2], inp.device
-
-        # get target
-
-        seq_start_pos = None
-
-        if one_step_autoregress:
-            target = x[:, None, 1:]
-        else:
-
-            batch_arange = arange(batch, device = device)
-            batch_arange = rearrange(batch_arange, 'b -> b 1 1')
-            seq_arange = arange(seq_len, device = device)
-            steps_arange = arange(steps, device = device) + 1
+        assert rollout_steps > 1
 
-            target_indices = einx.add('r, n -> r n', steps_arange, seq_arange)
+        steps = rollout_steps
 
-            target = x[batch_arange, target_indices] # rollout targets
-
-            seq_start_pos = torch.zeros(batch, device = device, dtype = torch.long)
+        device = x.device
 
         # assert inputs
 
@@ -348,53 +318,120 @@ class ContinuousAutoregressiveWrapper(Module):
             mask = einx.less('j, i -> i j', seq_arange, lens)
             kwargs['mask'] = mask
 
+        if not exists(lens):
+            batch, seq_len = x.shape[:2]
+            lens = torch.full((batch,), seq_len, device = device)
+
         # handle mask manually
 
         mask = kwargs.pop('mask', None)
 
-        has_mask = exists(mask)
+        # pick a random range for each batch sample and aligh the sequence to the right for rollout loss
 
-        # maybe rollout
+        valid_tokens_for_rollout = (lens - steps).clamp(min = 0)
+        valid_sample = valid_tokens_for_rollout > 0
+
+        x = x[valid_sample] # remove invalid sequence (lens less than rollout steps)
+
+        if exists(mask):
+            mask = mask[valid_sample]
+
+        batch = x.shape[0]
+        seq_start_pos = (torch.rand((batch,), device = device) * valid_tokens_for_rollout).floor().long()
+
+        batch_arange = torch.arange(batch, device = device)
+        batch_arange = rearrange(batch_arange, 'b -> b 1')
 
-        outputs = []
-        masks = []
+        # crop out sequence to use
 
-        for step_index in range(steps):
+        seq_end_pos = seq_start_pos + steps
+        max_end_pos = seq_end_pos.amax().item()
+        x = x[:, :max_end_pos]
 
-            step_mask = None
-            if has_mask:
-                step_mask = mask[:, step_index:(step_index + seq_len)]
-                masks.append(step_mask)
+        x = align_right(x, seq_end_pos)
 
-            # forward
+        # get the input
+
+        inp, targets = x[:, :-steps], x[:, -steps:]
 
-            out = self.net(inp, mask = step_mask, seq_start_pos = seq_start_pos, **kwargs)
+        # maybe rollout
 
-            outputs.append(out)
+        cache = None
+        preds = []
 
-            inp = out
+        for _ in range(steps):
 
-            if not one_step_autoregress:
-                seq_start_pos.sub_(1)
+            out, cache = self.net(
+                inp,
+                seq_start_pos = seq_start_pos,
+                return_intermediates = True,
+                **kwargs
+            )
 
-        # stack masks and predictions from rollouts
+            last_pred = out[..., -1:, :]
+
+            if self.probabilistic:
+                mean, var = last_pred
+                std = var.clamp(min = 1e-5).sqrt()
+                inp = torch.normal(mean, std)
+            else:
+                inp = last_pred
 
-        masks = stack(masks, dim = 1) if exists(mask) else None
+            preds.append(last_pred)
 
-        pred = stack(outputs, dim = 1)
+        # stack for predictions
+
+        preds = cat(preds, dim = 1)
 
         # loss
 
-        loss = self.loss_fn(pred, target)
+        loss = self.loss_fn(preds, targets)
+
+        return loss.mean()
+
+    def forward(
+        self,
+        x,
+        rollout_steps = 1, # they used 2 rollout steps in a successful world model paper https://ai.meta.com/vjepa/
+        **kwargs
+    ):
+        if rollout_steps > 1:
+            return self.forward_rollout(x, rollout_steps = rollout_steps, **kwargs)
+
+        inp, target = x[:, :-1], x[:, 1:]
+
+        assert 'prepend_embeds' not in kwargs
+
+        # lens
+
+        lens = kwargs.pop('lens', None)
+
+        if exists(lens):
+            assert 'mask' not in kwargs, 'either `mask` or `lens` passed in, but not both'
+            seq_len, device = inp.shape[1], inp.device
+            seq_arange = torch.arange(seq_len, device = device)
+            mask = einx.less('j, i -> i j', seq_arange, lens)
+
+            kwargs['mask'] = mask
+
+        # mask
+
+        mask = kwargs.get('mask', None)
+
+        if exists(mask) and mask.shape[1] == x.shape[1]:
+            mask = mask[:, :-1]
+            kwargs['mask'] = mask
+
+        out = self.net(inp, **kwargs)
 
-        # adjusting loss based on mask
+        loss = self.loss_fn(out, target)
 
-        if has_mask:
+        if exists(mask):
             assert loss.ndim > 1, 'loss should not be reduced if mask is passed in'
 
             if self.equal_loss_weight_batch:
-                loss = masked_mean(loss, masks)
+                loss = masked_mean(loss, mask)
             else:
-                loss = loss[masks]
+                loss = loss[mask]
 
         return loss.mean()

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

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

@@ -1,8 +1,8 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
 x_transformers/attend.py,sha256=xFsBtl7h7_qebPh7kE81BpmCWAjCgFpB9i_IHu_91es,17288
-x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
+x_transformers/autoregressive_wrapper.py,sha256=LW1gr3cFONDEPA_HHhaTE7mk-JWbaINuB1fc_DfbCqw,10791
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
-x_transformers/continuous.py,sha256=jy2wsQ3sS80Qwm_gnAmdAnzBfzLoWrGPacOTzU1Q6JM,11674
+x_transformers/continuous.py,sha256=uV2hLQOckeRsybqJy-0F8RhAyMPJlkVHmA7QqUJHG4g,12433
 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.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,,
+x_transformers-2.3.18.dist-info/METADATA,sha256=RKXNlO50fifu1Nas38iZRn6IJVDkv4Cen94XYVJlWg0,89897
+x_transformers-2.3.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.18.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.18.dist-info/RECORD,,