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

x_transformers/attend.py

@@ -69,7 +69,7 @@ def onnx_create_causal_mask(i, j, device):
 
 # main class
 
-class Attend(nn.Module):
+class Attend(Module):
     def __init__(
         self,
         *,
@@ -81,7 +81,8 @@ class Attend(nn.Module):
         scale = None,
         qk_norm = False,
         flash = False,
-        logit_softclamp_value = None,
+        softclamp_logits = False,
+        logit_softclamp_value = 30.,
         add_zero_kv = False,
         cope = None,
         onnxable = False,
@@ -123,10 +124,11 @@ class Attend(nn.Module):
 
         # soft clamp attention logit value
 
-        if exists(logit_softclamp_value):
+        if softclamp_logits:
             assert not flash, 'flash attention not compatible with logit softclamp value yet'
             assert logit_softclamp_value > 0.
 
+        self.softclamp_logits = softclamp_logits
         self.logit_softclamp_value = logit_softclamp_value
 
         # contextual positional encoding
@@ -308,6 +310,9 @@ class Attend(nn.Module):
         if exists(attn_bias):
             sim = sim + attn_bias
 
+        if self.softclamp_logits:
+            sim = softclamp(sim, self.logit_softclamp_value)
+
         i, j, dtype = *sim.shape[-2:], sim.dtype
 
         mask_value = -torch.finfo(sim.dtype).max
@@ -329,9 +334,6 @@ class Attend(nn.Module):
 
         pre_softmax_attn = sim.clone()
 
-        if exists(self.logit_softclamp_value):
-            sim = softclamp(sim, self.logit_softclamp_value)
-
         attn = self.attn_fn(sim, dim = -1)
         attn = attn.type(dtype)
 

x_transformers/x_transformers.py

@@ -884,7 +884,8 @@ class Attention(Module):
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
         cope_talking_heads = False,
-        logit_softclamp_value = None,
+        softclamp_logits = False,
+        logit_softclamp_value = 30.,
         onnxable = False
     ):
         super().__init__()
@@ -987,6 +988,7 @@ class Attention(Module):
             scale = qk_norm_scale if qk_norm else self.scale,
             add_zero_kv = add_zero_kv,
             flash = flash,
+            softclamp_logits = softclamp_logits,
             logit_softclamp_value = logit_softclamp_value,
             cope = cope,
             onnxable = onnxable
@@ -1265,7 +1267,8 @@ class AttentionLayers(Module):
         scale_residual_constant = 1.,
         shift_tokens = 0,
         sandwich_norm = False,
-        softclamp_output_value: float | None = None,
+        softclamp_output = False,
+        softclamp_output_value = 50.,
         resi_dual = False,
         resi_dual_scale = 1.,
         zero_init_branch_output = False,
@@ -1482,6 +1485,7 @@ class AttentionLayers(Module):
         # optional soft clamping just before the final norm
         # used in gemma 2
 
+        self.softclamp_output = softclamp_output
         self.softclamp_output_value = softclamp_output_value
 
         # whether it has post norm
@@ -1715,7 +1719,7 @@ class AttentionLayers(Module):
         if return_hiddens:
             layer_hiddens.append(x)
 
-        if exists(self.softclamp_output_value):
+        if self.softclamp_output:
             x = softclamp(x, self.softclamp_output_value)
 
         final_norm = self.final_norm

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

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

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

@@ -1,14 +1,14 @@
 x_transformers/__init__.py,sha256=8LQl-dNL6vj8VHRx5LMSOlRDTXQvYOuM21PDXz8WdiI,703
-x_transformers/attend.py,sha256=ap2QkD-bRadFE9ZFQP84Lo1P2DpLOXPam24Jq9ybpPY,10903
+x_transformers/attend.py,sha256=UWq0bElvJf-_j1N2QbJ2yg28xkWlhnOrLjMJt3If3ao,10956
 x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3nPqO6xp2FFCc,9619
 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=82gdvdwWWDWZVXTQRE6B4W-RQspg4SXhD8B6BV5pOo0,75789
+x_transformers/x_transformers.py,sha256=xPwzR3bd8BS_ChEcz0UxsNtx99u4UbP8jg1fFIRDGUw,75925
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-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,,
+x_transformers-1.31.6.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.31.6.dist-info/METADATA,sha256=O2MZXNuX-jqrAdpcxkIDH4J0T63t5nt9utI8FHrRIA0,661
+x_transformers-1.31.6.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
+x_transformers-1.31.6.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.31.6.dist-info/RECORD,,