x-transformers 2.3.17__tar.gz → 2.3.18__tar.gz

{x_transformers-2.3.17 → x_transformers-2.3.18}/PKG-INFO

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

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

{x_transformers-2.3.17 → x_transformers-2.3.18}/tests/test_x_transformers.py

@@ -876,9 +876,6 @@ def test_continuous(
     cache_kv,
     rollout_steps
 ):
-    if probabilistic and rollout_steps > 1:
-        pytest.skip()
-
     from x_transformers import (
         ContinuousTransformerWrapper,
         Decoder,

{x_transformers-2.3.17 → x_transformers-2.3.18}/x_transformers/continuous.py

@@ -298,7 +298,6 @@ class ContinuousAutoregressiveWrapper(Module):
         **kwargs
     ):
         assert rollout_steps > 1
-        assert not self.probabilistic, 'probabilistic not supported yet'
 
         steps = rollout_steps
 
@@ -369,8 +368,14 @@ class ContinuousAutoregressiveWrapper(Module):
                 **kwargs
             )
 
-            last_pred = out[:, -1:]
-            inp = last_pred
+            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
 
             preds.append(last_pred)