x-transformers 2.1.35__tar.gz → 2.1.36__tar.gz

{x_transformers-2.1.35 → x_transformers-2.1.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.35
+Version: 2.1.36
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.1.35 → x_transformers-2.1.36}/pyproject.toml

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

{x_transformers-2.1.35 → x_transformers-2.1.36}/x_transformers/x_transformers.py

@@ -864,14 +864,15 @@ class DynamicTanh(Module):
         self.gamma = nn.Parameter(torch.ones(dim))
         self.beta = nn.Parameter(torch.zeros(dim))
 
-        self.unit_offset = int(unit_offset)
+        self.pre_tanh_scale_offset = init_alpha if unit_offset else 0.
+        self.gamma_offset = float(unit_offset)
 
-        nn.init.constant_(self.pre_tanh_scale, 1. - float(unit_offset))
+        nn.init.constant_(self.pre_tanh_scale, 0 if unit_offset else init_alpha)
         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
+        pre_tanh_scale = self.pre_tanh_scale + self.pre_tanh_scale_offset
+        gamma = self.gamma + self.gamma_offset
         return (x * pre_tanh_scale).tanh() * gamma + self.beta
 
 # residual and residual gates