x-transformers 1.31.2__py3-none-any.whl → 1.31.4__py3-none-any.whl

x_transformers/x_transformers.py

@@ -566,7 +566,7 @@ class LayerNorm(Module):
     def __init__(
         self,
         dim,
-        unit_offset = 0.
+        unit_offset = False
     ):
         """
         bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less
@@ -576,9 +576,7 @@ class LayerNorm(Module):
 
         self.ln = nn.LayerNorm(dim, elementwise_affine = False)
         self.gamma = nn.Parameter(torch.ones(dim))
-        nn.init.constant_(self.gamma, 1. - unit_offset)
-
-        self.register_buffer('beta', torch.zeros(dim), persistent = False)
+        nn.init.constant_(self.gamma, 1. - float(unit_offset))
 
     def forward(self, x):
         normed = self.ln(x)
@@ -596,7 +594,6 @@ class AdaptiveLayerNorm(Module):
 
         self.ln = nn.LayerNorm(dim, elementwise_affine = False)
         self.to_gamma = nn.Linear(dim_condition, dim, bias = False)
-
         nn.init.zeros_(self.to_gamma.weight)
 
     def forward(self, x, *, condition):
@@ -608,14 +605,14 @@ class ScaleNorm(Module):
     def __init__(
         self,
         dim,
-        unit_offset = 0.
+        unit_offset = False
     ):
         super().__init__()
         self.unit_offset = unit_offset
         self.scale = dim ** 0.5
 
         self.g = nn.Parameter(torch.zeros(1))
-        nn.init.constant_(self.g, 1. - unit_offset)
+        nn.init.constant_(self.g, 1. - float(unit_offset))
 
     def forward(self, x):
         return F.normalize(x, dim = -1) * self.scale * (self.g + self.unit_offset)
@@ -624,14 +621,14 @@ class RMSNorm(Module):
     def __init__(
         self,
         dim,
-        unit_offset = 0.
+        unit_offset = False
     ):
         super().__init__()
         self.unit_offset = unit_offset
         self.scale = dim ** 0.5
 
         self.g = nn.Parameter(torch.zeros(dim))
-        nn.init.constant_(self.g, 1. - unit_offset)
+        nn.init.constant_(self.g, 1. - float(unit_offset))
 
     def forward(self, x):
         return F.normalize(x, dim = -1) * self.scale * (self.g + self.unit_offset)
@@ -738,14 +735,14 @@ class LayerScale(Module):
         fn: Module,
         dim,
         init_value = 0.,
-        unit_offset = 0.
+        unit_offset = False
     ):
         super().__init__()
         self.unit_offset = unit_offset
 
         self.fn = fn
         self.gamma = nn.Parameter(torch.zeros(dim))
-        nn.init.constant_(self.gamma, init_value - unit_offset)
+        nn.init.constant_(self.gamma, init_value - float(unit_offset))
 
     def forward(self, x, **kwargs):
         out = self.fn(x, **kwargs)
@@ -1370,8 +1367,8 @@ class AttentionLayers(Module):
         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 and betas, they can be subjected to weight decay safely
-            norm_fn = partial(norm_fn, unit_offset = 1.)
+            # researcher Ohad Rubin shares in a blog post by adding an offset to gammas, they can be subjected to weight decay safely
+            norm_fn = partial(norm_fn, unit_offset = True)
 
         self.norm_need_condition = norm_need_condition
         self.dim_condition = dim_condition
@@ -1403,8 +1400,8 @@ class AttentionLayers(Module):
 
         self.post_branch_fn_needs_condition = post_branch_fn_needs_condition
 
-        if not post_branch_fn_needs_condition and norm_add_unit_offset:
-            post_branch_fn = partial(post_branch_fn, unit_offset = 1.)
+        if exists(post_branch_fn) and not post_branch_fn_needs_condition and norm_add_unit_offset:
+            post_branch_fn = partial(post_branch_fn, unit_offset = True)
 
         # setup mlp for conditioning
 

{x_transformers-1.31.2.dist-info → x_transformers-1.31.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.31.2
+Version: 1.31.4
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.31.2.dist-info → x_transformers-1.31.4.dist-info}/RECORD

@@ -4,11 +4,11 @@ x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3n
 x_transformers/continuous.py,sha256=WO52n9lFAXv5-SGadi2cApGF8dkouN8QSTEOuC7erj8,6180
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=aeFAk19BRxkLOS-pUgtDk3HaMS6_dsuJtz2o_KkUkOU,75805
+x_transformers/x_transformers.py,sha256=82gdvdwWWDWZVXTQRE6B4W-RQspg4SXhD8B6BV5pOo0,75789
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.31.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.31.2.dist-info/METADATA,sha256=dxDZYzaZwwOCD-9I6ver-SSLlSG4diss-ovDo2H_9xE,661
-x_transformers-1.31.2.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
-x_transformers-1.31.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.31.2.dist-info/RECORD,,
+x_transformers-1.31.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.31.4.dist-info/METADATA,sha256=USK7uPjCATSS4URgfekTTmPhOtvHagHC02b-R-gxwME,661
+x_transformers-1.31.4.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
+x_transformers-1.31.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.31.4.dist-info/RECORD,,