x-transformers 2.1.29__py3-none-any.whl → 2.1.30__py3-none-any.whl

x_transformers/x_transformers.py

@@ -9,7 +9,7 @@ from packaging import version
 import torch
 from torch.amp import autocast
 import torch.nn.functional as F
-from torch import nn, einsum, Tensor, cat, stack, arange, is_tensor
+from torch import nn, einsum, tensor, Tensor, cat, stack, arange, is_tensor
 from torch.utils._pytree import tree_flatten, tree_unflatten
 from torch.nn import Module, ModuleList, ModuleDict
 
@@ -266,7 +266,6 @@ class TokenEmbedding(Module):
             return
         nn.init.kaiming_normal_(self.emb.weight)
 
-
 # positional embeddings
 
 class AbsolutePositionalEmbedding(Module):
@@ -849,6 +848,31 @@ class MultiheadRMSNorm(Module):
     def forward(self, x):
         return self.rmsnorm(x) * (self.gamma + 1.)
 
+class DynamicTanh(Module):
+    """ https://arxiv.org/abs/2503.10622 """
+    def __init__(
+        self,
+        init_alpha = 1.,
+        gamma = 1.,
+        beta = 0.,
+        unit_offset = False
+    ):
+        super().__init__()
+        self.pre_tanh_scale = nn.Parameter(tensor(init_alpha))
+
+        self.gamma = nn.Parameter(tensor(init_alpha))
+        self.beta = nn.Parameter(tensor(init_alpha))
+
+        self.unit_offset = int(unit_offset)
+
+        nn.init.constant_(self.pre_tanh_scale, 1. - float(unit_offset))
+        nn.init.constant_(self.gamma, 1. - float(unit_offset))
+
+    def forward(self, x):
+        pre_tanh_scale = self.pre_tanh_scale + self.unit_offset
+        gamma = self.gamma + self.unit_offset
+        return (x * pre_tanh_scale).tanh() * gamma + self.beta
+
 # residual and residual gates
 
 class Residual(Module):
@@ -1863,6 +1887,8 @@ class AttentionLayers(Module):
         only_cross = False,
         use_scalenorm = False,
         use_rmsnorm = False,
+        use_dynamic_tanh = False,
+        dynamic_tanh_init_alpha = 1.,
         use_simple_rmsnorm = False,
         use_adaptive_layernorm = False,
         use_adaptive_rmsnorm = False,
@@ -2012,8 +2038,9 @@ class AttentionLayers(Module):
 
         # determine norm
 
-        assert at_most_one_of(use_scalenorm, use_rmsnorm, use_simple_rmsnorm, use_adaptive_layernorm, use_adaptive_rmsnorm), 'you can only use either scalenorm, rmsnorm, adaptive layernorm, adaptive rmsnorm, or simple rmsnorm'
+        assert at_most_one_of(use_scalenorm, use_rmsnorm, use_dynamic_tanh, use_simple_rmsnorm, use_adaptive_layernorm, use_adaptive_rmsnorm), 'you can only use either scalenorm, rmsnorm, adaptive layernorm, adaptive rmsnorm, or simple rmsnorm'
 
+        norm_fn = None
         norm_need_condition = False
         dim_condition = default(dim_condition, dim)
         dim_condition_mult = 1
@@ -2027,6 +2054,8 @@ class AttentionLayers(Module):
             norm_class = RMSNorm
         elif use_simple_rmsnorm:
             norm_class = SimpleRMSNorm
+        elif use_dynamic_tanh:
+            norm_fn = partial(DynamicTanh, init_alpha = dynamic_tanh_init_alpha)
         elif use_adaptive_layernorm:
             norm_need_condition = True
             norm_class = partial(AdaptiveLayerNorm, dim_condition = dim_condition * dim_condition_mult)
@@ -2036,7 +2065,8 @@ class AttentionLayers(Module):
         else:
             norm_class = LayerNorm
 
-        norm_fn = partial(norm_class, dim)
+        if not exists(norm_fn):
+            norm_fn = partial(norm_class, dim)
 
         if not norm_need_condition and norm_add_unit_offset:
             # researcher Ohad Rubin shares in a blog post by adding an offset to gammas, they can be subjected to weight decay safely

{x_transformers-2.1.29.dist-info → x_transformers-2.1.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.29
+Version: 2.1.30
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2455,4 +2455,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Zhu2025TransformersWN,
+    title   = {Transformers without Normalization},
+    author  = {Jiachen Zhu and Xinlei Chen and Kaiming He and Yann LeCun and Zhuang Liu},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:276961218}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.1.29.dist-info → x_transformers-2.1.30.dist-info}/RECORD

@@ -7,10 +7,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=fqgtIs6__JpLWMnJa8AY5OW3AJ2GR1B5p-9TsWdiOIU,110425
+x_transformers/x_transformers.py,sha256=a7k6tR9H1kCRX44PP0N9nMMb3V1_cIgFweTBK84VtEk,111476
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.1.29.dist-info/METADATA,sha256=CI6GLna-OqlmDEjv8sP0CcfI7SNCAbL-nQCm2sQqdbc,87875
-x_transformers-2.1.29.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.29.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.29.dist-info/RECORD,,
+x_transformers-2.1.30.dist-info/METADATA,sha256=IdKgXNQf9aTZ_JiOYhc3q1J44ITmwBO-VRPRDaZtnEU,88161
+x_transformers-2.1.30.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.30.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.30.dist-info/RECORD,,