x-transformers 1.30.1__py3-none-any.whl → 1.30.3__py3-none-any.whl

x_transformers/attend.py

@@ -37,6 +37,9 @@ def default(val, d):
 def compact(arr):
     return [*filter(exists, arr)]
 
+def softclamp(t, value):
+    return (t / value).tanh() * value
+
 def once(fn):
     called = False
     @wraps(fn)
@@ -76,6 +79,7 @@ class Attend(nn.Module):
         scale = None,
         qk_norm = False,
         flash = False,
+        logit_softclamp_value = None,
         add_zero_kv = False,
         onnxable = False,
         sdp_kwargs: dict = dict(
@@ -114,6 +118,14 @@ class Attend(nn.Module):
 
         self.add_zero_kv = add_zero_kv
 
+        # soft clamp attention logit value
+
+        if exists(logit_softclamp_value):
+            assert not flash, 'flash attention not compatible with logit softclamp value yet'
+            assert logit_softclamp_value > 0.
+
+        self.logit_softclamp_value = logit_softclamp_value
+
         # flash attention
 
         self.flash = flash
@@ -276,38 +288,41 @@ class Attend(nn.Module):
 
         kv_einsum_eq = 'b j d' if k.ndim == 3 else 'b h j d'
 
-        dots = einsum(f'b h i d, {kv_einsum_eq} -> b h i j', q, k) * scale
+        sim = einsum(f'b h i d, {kv_einsum_eq} -> b h i j', q, k) * scale
 
         if exists(prev_attn):
-            dots = dots + prev_attn
+            sim = sim + prev_attn
 
-        qk_similarities = dots.clone()
+        qk_similarities = sim.clone()
 
         if self.talking_heads:
-            dots = self.pre_softmax_talking_heads(dots)
+            sim = self.pre_softmax_talking_heads(sim)
 
         if exists(attn_bias):
-            dots = dots + attn_bias
+            sim = sim + attn_bias
 
-        i, j, dtype = *dots.shape[-2:], dots.dtype
+        i, j, dtype = *sim.shape[-2:], sim.dtype
 
-        mask_value = -torch.finfo(dots.dtype).max
+        mask_value = -torch.finfo(sim.dtype).max
 
         if exists(self.sparse_topk) and self.sparse_topk < j:
-            top_values, _ = dots.topk(self.sparse_topk, dim = -1)
-            sparse_topk_mask = dots < top_values[..., -1:]
+            top_values, _ = sim.topk(self.sparse_topk, dim = -1)
+            sparse_topk_mask = sim < top_values[..., -1:]
             mask = (mask & sparse_topk_mask) if exists(mask) else sparse_topk_mask
 
         if exists(mask):
-            dots = dots.masked_fill(~mask, mask_value)
+            sim = sim.masked_fill(~mask, mask_value)
 
         if causal:
             causal_mask = self.create_causal_mask(i, j, device = device)
-            dots = dots.masked_fill(causal_mask, mask_value)
+            sim = sim.masked_fill(causal_mask, mask_value)
+
+        pre_softmax_attn = sim.clone()
 
-        pre_softmax_attn = dots.clone()
+        if exists(self.logit_softclamp_value):
+            sim = softclamp(sim, self.logit_softclamp_value)
 
-        attn = self.attn_fn(dots, dim = -1)
+        attn = self.attn_fn(sim, dim = -1)
         attn = attn.type(dtype)
 
         post_softmax_attn = attn.clone()

x_transformers/x_transformers.py

@@ -468,7 +468,8 @@ def rotate_half(x):
 
 @autocast(enabled = False)
 def apply_rotary_pos_emb(t, freqs, scale = 1):
-    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
+    rot_dim, seq_len, orig_dtype = freqs.shape[-1], t.shape[-2], t.dtype
+
     freqs = freqs[-seq_len:, :]
     scale = scale[-seq_len:, :] if isinstance(scale, torch.Tensor) else scale
 
@@ -478,7 +479,9 @@ def apply_rotary_pos_emb(t, freqs, scale = 1):
     # partial rotary embeddings, Wang et al. GPT-J
     t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
     t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
-    return torch.cat((t, t_unrotated), dim = -1)
+    out = torch.cat((t, t_unrotated), dim = -1)
+
+    return out.type(orig_dtype)
 
 # norms
 
@@ -719,6 +722,7 @@ class Attention(Module):
         tensor_product = False,      # https://arxiv.org/abs/2208.06061
         add_zero_kv = False,         # same as add_zero_attn in pytorch
         rotary_embed_values = False,
+        logit_softclamp_value = None,
         onnxable = False
     ):
         super().__init__()
@@ -798,6 +802,7 @@ class Attention(Module):
             scale = qk_norm_scale if qk_norm else self.scale,
             add_zero_kv = add_zero_kv,
             flash = flash,
+            logit_softclamp_value = logit_softclamp_value,
             onnxable = onnxable
         )
 

{x_transformers-1.30.1.dist-info → x_transformers-1.30.3.dist-info}/METADATA

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

{x_transformers-1.30.1.dist-info → x_transformers-1.30.3.dist-info}/RECORD

@@ -1,14 +1,14 @@
 x_transformers/__init__.py,sha256=8LQl-dNL6vj8VHRx5LMSOlRDTXQvYOuM21PDXz8WdiI,703
-x_transformers/attend.py,sha256=Y9eE26I7BM8rGveabhiRhzw_xq9TY61Sp10QC1hX2O8,10192
+x_transformers/attend.py,sha256=2SPHjXS_QAAZt04lHWGtdOypTExmo3BrbFhgcIQTk-Y,10671
 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=EEfqwI-NANzrQf10Tc_bRSdjWOIEJdhxOfzeKY4osyI,66137
+x_transformers/x_transformers.py,sha256=BQypGJAoqXrAe_ek95wUcXdSQdAWjvw5mEti-H1JxcI,66288
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.30.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.30.1.dist-info/METADATA,sha256=gkmRLAvk0l9_vkrTVBIWLnFq_cEtCYrI8oI3B07d9B8,661
-x_transformers-1.30.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-x_transformers-1.30.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.30.1.dist-info/RECORD,,
+x_transformers-1.30.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.30.3.dist-info/METADATA,sha256=FdVMtNhhggibbG3fJfCfGdRUAp6fLkTyG8KzDUq_r1Y,661
+x_transformers-1.30.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+x_transformers-1.30.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.30.3.dist-info/RECORD,,