x-transformers 2.3.17__py3-none-any.whl → 2.3.19__py3-none-any.whl

x_transformers/attend.py

@@ -25,6 +25,7 @@ class Intermediates:
     values:             Tensor | None = None
     cached_kv:          Tuple[Tensor, Tensor] | None = None
     layer_type:         str | None = None
+    hybrid_hidden:      Tensor | None = None
 
     def to_tuple(self):
         return (self.qk_similarities, self.pre_softmax_attn, self.post_softmax_attn)

x_transformers/continuous.py

@@ -32,6 +32,15 @@ def default(val, d):
         return val
     return d() if not isinstance(d, Module) and callable(d) else d
 
+def sample_from_mean_variance(
+    mean,
+    variance,
+    eps = 1e-5,
+    temperature = 1.
+):
+    std = variance.clamp(min = eps).sqrt()
+    return torch.normal(mean, std * temperature)
+
 def masked_mean(t, mask):
     t = einx.where('b n, b n d, -> b n d', mask, t, 0.)
 
@@ -274,9 +283,7 @@ class ContinuousAutoregressiveWrapper(Module):
 
             if self.probabilistic:
                 mean, var = last_output
-                stddev = var.clamp(min = 1e-5).sqrt()
-
-                last_output = torch.normal(mean, stddev * temperature)
+                last_output = sample_from_mean_variance(mean, var, temperature = temperature)
 
             out = cat((out, last_output), dim = -2)
 
@@ -298,7 +305,6 @@ class ContinuousAutoregressiveWrapper(Module):
         **kwargs
     ):
         assert rollout_steps > 1
-        assert not self.probabilistic, 'probabilistic not supported yet'
 
         steps = rollout_steps
 
@@ -369,8 +375,13 @@ class ContinuousAutoregressiveWrapper(Module):
                 **kwargs
             )
 
-            last_pred = out[:, -1:]
-            inp = last_pred
+            last_pred = out[..., -1:, :]
+
+            if self.probabilistic:
+                mean, var = last_pred
+                inp = sample_from_mean_variance(mean, var)
+            else:
+                inp = last_pred
 
             preds.append(last_pred)
 

x_transformers/x_transformers.py

@@ -1079,10 +1079,11 @@ class FoldAxially(Module):
     def forward(
         self,
         x,
+        *args,
         **kwargs
     ):
         if self.axial_dim == 1:
-            return self.fn(x, **kwargs)
+            return self.fn(x, *args, **kwargs)
 
         seq_len, axial_dim = x.shape[1], self.axial_dim
 
@@ -1091,7 +1092,7 @@ class FoldAxially(Module):
 
         x = rearrange(x, 'b (n axial_dim) ... -> (b axial_dim) n ...', axial_dim = axial_dim)
 
-        out = self.fn(x, **kwargs)
+        out = self.fn(x, *args, **kwargs)
 
         (out, *rest_out), tree_spec = tree_flatten(out)
 
@@ -1857,9 +1858,17 @@ class Attention(Module):
             if not self.causal and exists(self.hybrid_mask_kwarg):
                 hybrid_forward_kwargs = {self.hybrid_mask_kwarg: mask}
 
+            # handle maybe hybrid cache
+
+            hybrid_forward_args = ()
+
+            if exists(cache) and exists(cache.hybrid_hidden):
+                hybrid_hiddens = cache.hybrid_hidden
+                hybrid_forward_args = (hybrid_hiddens,)
+
             # hybrid forward
 
-            hybrid_outputs = self.hybrid_module(x, **hybrid_forward_kwargs)
+            hybrid_outputs = self.hybrid_module(x, *hybrid_forward_args, **hybrid_forward_kwargs)
 
             # handle hybrid out
 
@@ -1870,6 +1879,10 @@ class Attention(Module):
             if hybrid_out.ndim == 3:
                 hybrid_out = rearrange(hybrid_out, 'b n (h d) -> b h n d', h = h)
 
+            if len(rest_hybrid_outs) > 0:
+                hybrid_hidden = first(rest_hybrid_outs)
+                intermediates.hybrid_hidden = hybrid_hidden
+
             out_norm, hybrid_out_norm = self.hybrid_norms
 
             out = out_norm(out)

{x_transformers-2.3.17.dist-info → x_transformers-2.3.19.dist-info}/METADATA

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

@@ -1,17 +1,17 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
-x_transformers/attend.py,sha256=xFsBtl7h7_qebPh7kE81BpmCWAjCgFpB9i_IHu_91es,17288
+x_transformers/attend.py,sha256=fXMuwHuBAFB4f4_U6j5_uVeK7N4cV0PDd6UTqtkjKKM,17333
 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=QidhbSgBhYDpAp7FaryxrNJxWXp0-pJkeygguTpVp4k,12308
+x_transformers/continuous.py,sha256=CHta8vizKl85n220fv5278fwjSU-vrN_FBy-m831_go,12551
 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
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=ZfOXrZSiy2jlZ8wVmDdMTLW4hAY_qfmPQHW9t2ABxbo,114097
+x_transformers/x_transformers.py,sha256=s398YQ9JtXc5n34g9qaYnUqaTVLGfRvz0GLg3sEMHLI,114558
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.17.dist-info/METADATA,sha256=u_DVXsX7KsVhnfMhpV-3KV6KGNHddmzW_SNG9om557s,89897
-x_transformers-2.3.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.17.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.17.dist-info/RECORD,,
+x_transformers-2.3.19.dist-info/METADATA,sha256=Vn-U7mDaP7H-w-RF5YO3C5n9M5PvnDVKqFJwL3vFV0s,89897
+x_transformers-2.3.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.19.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.19.dist-info/RECORD,,