x-transformers 2.1.28__tar.gz → 2.1.30__tar.gz

{x_transformers-2.1.28 → x_transformers-2.1.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.28
+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.28 → x_transformers-2.1.30}/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.28 → x_transformers-2.1.30}/pyproject.toml

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

{x_transformers-2.1.28 → x_transformers-2.1.30}/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.28 → x_transformers-2.1.30}/x_transformers/belief_state_wrapper.py

@@ -258,6 +258,8 @@ class BeliefStateWrapper(Module):
 
         # handle variable length sequences
 
+        seq_for_labels = seq
+
         if exists(lens):
             mask = einx.less('j, i -> i j', arange(seq_len, device = device), lens)
             seq_for_labels = torch.where(mask, seq, -1)

{x_transformers-2.1.28 → x_transformers-2.1.30}/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