x-transformers 2.3.2__tar.gz → 2.3.4__tar.gz

{x_transformers-2.3.2 → x_transformers-2.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.3.2
+Version: 2.3.4
 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.2 → x_transformers-2.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.3.2"
+version = "2.3.4"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.3.2 → x_transformers-2.3.4}/x_transformers/continuous.py

@@ -123,6 +123,7 @@ class ContinuousTransformerWrapper(Module):
         return_intermediates = False,
         return_mems = False,
         mask = None,
+        lens = None,
         return_attn = False,
         mems = None,
         mem_masks = None,
@@ -133,6 +134,16 @@ class ContinuousTransformerWrapper(Module):
     ):
         batch, seq, orig_mask, device = *x.shape[:2], mask, x.device
 
+        # maybe seq lengths passed in
+
+        if exists(lens):
+            assert not exists(mask), 'either `mask` or `lens` passed in, but not both'
+            seq_arange = torch.arange(seq, device = device)
+
+            mask = einx.less('j, i -> i j', seq_arange, lens)
+
+        # project in + positional embedding
+
         x = self.project_in(x)
         x = x + self.pos_emb(x, pos = pos)
 
@@ -282,7 +293,22 @@ class ContinuousAutoregressiveWrapper(Module):
 
         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