x-transformers 2.1.29__tar.gz → 2.1.31__tar.gz

{x_transformers-2.1.29 → x_transformers-2.1.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.29
+Version: 2.1.31
 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 → x_transformers-2.1.31}/README.md

@@ -2407,4 +2407,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 → x_transformers-2.1.31}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.1.29"
+version = "2.1.31"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.1.29 → x_transformers-2.1.31}/tests/test_x_transformers.py

@@ -697,10 +697,12 @@ def test_lime(
 @pytest.mark.parametrize('backward_ar_loss_weight', (1., 0.5))
 @pytest.mark.parametrize('goal_suffix', (False, True))
 @pytest.mark.parametrize('pred_distance', (False, True))
+@pytest.mark.parametrize('variable_len', (False, True))
 def test_belief_state_wrapper(
     backward_ar_loss_weight,
     goal_suffix,
-    pred_distance
+    pred_distance,
+    variable_len
 ):
     from x_transformers.belief_state_wrapper import BeliefStateWrapper
 
@@ -735,7 +737,12 @@ def test_belief_state_wrapper(
 
     seq = torch.randint(0, 20000, (2, 16))
 
-    loss = model(seq) # backwards happen automatically
+    lens = None
+
+    if variable_len:
+        lens = torch.randint(4, 16, (2,))
+
+    loss = model(seq, lens = lens) # backwards happen automatically
 
     suffix = None
     if goal_suffix:
@@ -743,3 +750,20 @@ def test_belief_state_wrapper(
 
     sampled = model.generate_with_suffix_cond(seq[:, :1], 16, suffix = suffix)
     assert sampled.shape == (2, 16)
+
+def test_dynamic_tanh():
+    model = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 128,
+            depth = 6,
+            heads = 8,
+            use_dynamic_tanh = True,
+            dynamic_tanh_init_alpha = 1.5
+        )
+    )
+
+    x = torch.randint(0, 20000, (2, 1024))
+
+    model(x)

{x_transformers-2.1.29 → x_transformers-2.1.31}/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,32 @@ 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,
+        dim,
+        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(torch.ones(dim))
+        self.beta = nn.Parameter(torch.zeros(dim))
+
+        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 +1888,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,7 +2039,7 @@ 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_need_condition = False
         dim_condition = default(dim_condition, dim)
@@ -2027,6 +2054,8 @@ class AttentionLayers(Module):
             norm_class = RMSNorm
         elif use_simple_rmsnorm:
             norm_class = SimpleRMSNorm
+        elif use_dynamic_tanh:
+            norm_class = 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)