x-transformers 2.3.2__py3-none-any.whl → 2.3.4__py3-none-any.whl

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

{x_transformers-2.3.2.dist-info → x_transformers-2.3.4.dist-info}/METADATA

@@ -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.dist-info → x_transformers-2.3.4.dist-info}/RECORD

@@ -2,7 +2,7 @@ x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,98
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
-x_transformers/continuous.py,sha256=yAE8hLyusrEd-12mkgLASDL-cFgpZQf32s93FKfez7o,8674
+x_transformers/continuous.py,sha256=TwqPXZG14JGzfK3iK4bdQo1u0u2KKNLRcRjh54Jc0z8,9422
 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=MF91aJGr2DOjIGe57uqwgyNxCExBg_tI9z7usAJMxOM,112401
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.3.2.dist-info/METADATA,sha256=8m6hpJlMKesI-SLxth_9z0VYIHUh7bTWsJ9Am5OSni4,88686
-x_transformers-2.3.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.2.dist-info/RECORD,,
+x_transformers-2.3.4.dist-info/METADATA,sha256=dRdw8JjntClC_behPfg-Upf0bGtuhn1UKklkXRrLCmU,88686
+x_transformers-2.3.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.4.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.4.dist-info/RECORD,,