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

x_transformers/gpt_vae.py

@@ -46,25 +46,29 @@ class GPTVAE(Module):
         vae_kl_loss_weight = 1.,
         latents_dropout_prob = 0.5, # what percentage of the time to dropout the latents completely
         pad_id = -1,
+        encoder: Module | None = None,
         **kwargs
     ):
         super().__init__()
         dim_latent = default(dim_latent, dim)
 
-        self.encoder = TransformerWrapper(
-            num_tokens = num_tokens,
-            max_seq_len = max_seq_len + 1,
-            return_only_embed = True,
-            average_pool_embed = True,
-            attn_layers = Encoder(
-                dim = dim,
-                depth = enc_depth,
-                attn_dim_head = attn_dim_head,
-                heads = heads,
-                **kwargs,
-                **enc_kwargs
-            ),
-        )
+        if not exists(encoder):
+            encoder = TransformerWrapper(
+                num_tokens = num_tokens,
+                max_seq_len = max_seq_len + 1,
+                return_only_embed = True,
+                average_pool_embed = True,
+                attn_layers = Encoder(
+                    dim = dim,
+                    depth = enc_depth,
+                    attn_dim_head = attn_dim_head,
+                    heads = heads,
+                    **kwargs,
+                    **enc_kwargs
+                ),
+            )
+
+        self.encoder = encoder
 
         self.to_latent_mean_log_variance = nn.Sequential(
             nn.Linear(dim, dim_latent * 2),

x_transformers/x_transformers.py

@@ -4,6 +4,11 @@ from typing import Callable
 import math
 from copy import deepcopy
 from random import random, randrange
+from functools import partial, wraps
+from itertools import chain
+from collections import namedtuple
+from contextlib import nullcontext
+from dataclasses import dataclass
 from packaging import version
 
 import torch
@@ -13,11 +18,6 @@ from torch import nn, einsum, tensor, Tensor, cat, stack, arange, is_tensor
 from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from torch.nn import Module, ModuleList, ModuleDict
 
-from functools import partial, wraps
-from collections import namedtuple
-from contextlib import nullcontext
-from dataclasses import dataclass
-
 from loguru import logger
 
 from x_transformers.attend import Attend, Intermediates
@@ -1279,6 +1279,17 @@ class FeedForward(Module):
         if zero_init_output:
             init_zero_(proj_out)
 
+    def muon_parameters(self):
+        weights = []
+
+        for m in self.modules():
+            if not isinstance(m, nn.Linear):
+                continue
+
+            weights.append(m.weight)
+
+        return weights
+
     def forward(
         self,
         x,
@@ -1644,6 +1655,9 @@ class Attention(Module):
         q_weight.mul_(qk_weight_scale)
         k_weight.mul_(qk_weight_scale)
 
+    def muon_parameters(self):
+        return chain(self.to_v.parameters(), self.to_out.parameters())
+
     def forward(
         self,
         x,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.8.1
+Version: 2.8.3
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2552,4 +2552,25 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{jordan2024muon,
+  author    = {Keller Jordan and Yuchen Jin and Vlado Boza and Jiacheng You and Franz Cesista and Laker Newhouse and Jeremy Bernstein},
+  title     = {Muon: An optimizer for hidden layers in neural networks},
+  year      = {2024},
+  url       = {https://kellerjordan.github.io/posts/muon/}
+}
+```
+
+```bibtex
+@misc{wang2025muonoutperformsadamtailend,
+    title   = {Muon Outperforms Adam in Tail-End Associative Memory Learning}, 
+    author  = {Shuche Wang and Fengzhuo Zhang and Jiaxiang Li and Cunxiao Du and Chao Du and Tianyu Pang and Zhuoran Yang and Mingyi Hong and Vincent Y. F. Tan},
+    year    = {2025},
+    eprint  = {2509.26030},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2509.26030}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

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

@@ -5,15 +5,15 @@ x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTN
 x_transformers/continuous.py,sha256=WwpQCjyVY4PtuEAOFY68zqgklbF9I7AL5w6874YlDe8,13249
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
-x_transformers/gpt_vae.py,sha256=Q2pzQ6iXRnP2Bfa6g-fs4US-JTouXB5-MfKw3sTwWmU,5561
+x_transformers/gpt_vae.py,sha256=myYSgcx66V0M4zeEGKyhY1P2HlPDHcezhaZEoo_uMdo,5715
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 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=odnCZAKZKrQLXmpaWhiPVB5elGjt8kerDbO3-yeC-60,124764
+x_transformers/x_transformers.py,sha256=gnmhtxPdmVQTd59MFXcGSm9HCKH9jv1fTBBYWAu5qaI,125113
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.8.1.dist-info/METADATA,sha256=_PnvoOSFJAgrpEfpNNljxdeYQ3BhDYJvVOp7yjaF-iM,94136
-x_transformers-2.8.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.8.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.8.1.dist-info/RECORD,,
+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,,