x-transformers 2.8.3__py3-none-any.whl → 2.9.0__py3-none-any.whl

x_transformers/x_transformers.py

@@ -2493,6 +2493,17 @@ class AttentionLayers(Module):
         for attn_layer, attn_inter in zip(attn_layers, attn_intermeds):
             attn_layer.qk_clip_(attn_inter, tau = tau)
 
+    def muon_parameters(self):
+        params = []
+
+        for m in self.modules():
+            if not isinstance(m, (Attention, FeedForward)):
+                continue
+
+            params.extend(list(m.muon_parameters()))
+
+        return params
+
     def forward(
         self,
         x,
@@ -3230,6 +3241,9 @@ class TransformerWrapper(Module):
     ):
         self.attn_layers.attn_qk_clip_(intermediates, tau = tau)
 
+    def muon_parameters(self):
+        return self.attn_layers.muon_parameters()
+
     def forward(
         self,
         x,

{x_transformers-2.8.3.dist-info → x_transformers-2.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.8.3
+Version: 2.9.0
 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: adam-atan2-pytorch>=0.2.2; extra == 'examples'
 Requires-Dist: lion-pytorch; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'
 Provides-Extra: test

{x_transformers-2.8.3.dist-info → x_transformers-2.9.0.dist-info}/RECORD

@@ -10,10 +10,10 @@ x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=gnmhtxPdmVQTd59MFXcGSm9HCKH9jv1fTBBYWAu5qaI,125113
+x_transformers/x_transformers.py,sha256=xIaM5LGimNmDThr9WK0ZF5vfm3sSp0OMll0aRIqjymQ,125443
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.8.3.dist-info/METADATA,sha256=vB7jRRZOX58zB9QhBagiQ3u61t6Xd6XMzWwnDngroVw,94924
-x_transformers-2.8.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.8.3.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.8.3.dist-info/RECORD,,
+x_transformers-2.9.0.dist-info/METADATA,sha256=tQ93i_TSfQIf9OViWNJj57ezY0Slred1NblIyf4UIs0,94986
+x_transformers-2.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.9.0.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.9.0.dist-info/RECORD,,