x-transformers 2.0.1__py3-none-any.whl → 2.0.2__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1204,6 +1204,7 @@ class Attention(Module):
         hybrid_module: Module | None = None,
         hybrid_mask_kwarg: str | None = None,
         hybrid_fold_axial_dim: int | None = None,
+        hybrid_learned_mix = False,
         one_kv_head = False,
         kv_heads = None,
         value_dim_head = None,
@@ -1446,7 +1447,7 @@ class Attention(Module):
 
         if exists(hybrid_module) and exists(hybrid_fold_axial_dim):
             hybrid_module = FoldAxially(axial_dim = hybrid_fold_axial_dim, fn = hybrid_module)
-            hybrid_mix = LinearNoBias(dim, heads)
+            hybrid_mix = LinearNoBias(dim, heads) if hybrid_learned_mix else None
 
             hybrid_norms = ModuleList([
                 MultiheadRMSNorm(dim_head, heads = heads),
@@ -1779,7 +1780,12 @@ class Attention(Module):
             out = out_norm(out)
             hybrid_out = hybrid_out_norm(hybrid_out)
 
-            out = 0.5 * (out + hybrid_out)
+            if exists(self.hybrid_mix):
+                mix = self.hybrid_mix(x)
+                mix = rearrange(mix, 'b n h -> b h n 1')
+                out = out.lerp(hybrid_out, mix.sigmoid())
+            else:
+                out = 0.5 * (out + hybrid_out)
 
         # merge heads
 

{x_transformers-2.0.1.dist-info → x_transformers-2.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.0.1
+Version: 2.0.2
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -40,6 +40,7 @@ Requires-Dist: loguru
 Requires-Dist: packaging>=21.0
 Requires-Dist: torch>=2.0
 Provides-Extra: examples
+Requires-Dist: lion-pytorch; extra == 'examples'
 Requires-Dist: torchvision; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'
 Provides-Extra: test

{x_transformers-2.0.1.dist-info → x_transformers-2.0.2.dist-info}/RECORD

@@ -6,10 +6,10 @@ x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 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=e-iIpcUJzDKU-_s4CVP9f7HoM2c11RH35-TH4G8pAYg,107251
+x_transformers/x_transformers.py,sha256=1s8KCSfHXMN9TKLFdS-RzzCskBDkh4CuBk2_XRb6IXk,107537
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.0.1.dist-info/METADATA,sha256=wygqpzU7ne5w7wj1GhPUEEuVROu0qIP0eDpjfaxlBe0,86457
-x_transformers-2.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.0.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.0.1.dist-info/RECORD,,
+x_transformers-2.0.2.dist-info/METADATA,sha256=tNdI3H2S4HnnGK1hPY3l94FoXH3SB9vGAb55pcah6Yw,86506
+x_transformers-2.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.0.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.0.2.dist-info/RECORD,,