x-transformers 1.37.9__py3-none-any.whl → 1.37.10__py3-none-any.whl

x_transformers/attend.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Tuple
+from typing import Tuple, Callable
 
 import torch
 from torch.nn import Module
@@ -104,6 +104,7 @@ class Attend(Module):
         qk_norm = False,
         l2_distance = False,
         sigmoid = False,
+        custom_attn_fn: Callable | None = None,
         flash = False,
         softclamp_logits = False,
         logit_softclamp_value = 50.,
@@ -132,7 +133,9 @@ class Attend(Module):
 
         self.sigmoid = sigmoid
 
-        if not sigmoid:
+        if exists(custom_attn_fn):
+            self.attn_fn = custom_attn_fn
+        elif not sigmoid:
             softmax_fn = partial(F.softmax, dim = -1)
             self.attn_fn = partial(softmax_fn, dtype = torch.float32) if not qk_norm else softmax_fn
         else:
@@ -404,11 +407,6 @@ class Attend(Module):
 
         mask_value = -torch.finfo(sim.dtype).max
 
-        if exists(self.sparse_topk) and self.sparse_topk < j:
-            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):
             sim = sim.masked_fill(~mask, mask_value)
 
@@ -416,6 +414,11 @@ class Attend(Module):
             causal_mask = self.create_causal_mask(i, j, device = device)
             sim = sim.masked_fill(causal_mask, mask_value)
 
+        if exists(self.sparse_topk):
+            top_values, _ = sim.topk(self.sparse_topk, dim = -1)
+            sparse_topk_mask = (sim >= top_values[..., -1:]) & (sim > mask_value)
+            sim = sim.masked_fill(~sparse_topk_mask, mask_value)
+
         row_is_entirely_masked = None
 
         if exists(mask):

x_transformers/x_transformers.py

@@ -925,6 +925,7 @@ class Attention(Module):
         qk_norm_dim_scale = False,
         l2_distance = False,
         sigmoid = False,
+        custom_attn_fn: Callable | None = None,
         one_kv_head = False,
         kv_heads = None,
         shared_kv = False,
@@ -1041,6 +1042,7 @@ class Attention(Module):
             scale = qk_norm_scale if qk_norm else self.scale,
             l2_distance = l2_distance,
             sigmoid = sigmoid,
+            custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,
             flash = flash,
             softclamp_logits = softclamp_logits,

{x_transformers-1.37.9.dist-info → x_transformers-1.37.10.dist-info}/METADATA

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

{x_transformers-1.37.9.dist-info → x_transformers-1.37.10.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=EtaTN6ahgRlFLkwfHA31RNL_bQyAHwhNBpGU1NIHJ-c,13894
+x_transformers/attend.py,sha256=MJmMLIt0rHx-JNNsc2auUsCjsB-69NewufaRV32ADmA,14012
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=Y77_ZPWSKTJ-oYk4bHjhwMEkgoMaq_LyxcmCkvOPZ9g,83808
+x_transformers/x_transformers.py,sha256=q6I6rvyYUWLgwtKOxPwF12UL1HzcIlauI8YrM8gvZac,83901
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.37.9.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.37.9.dist-info/METADATA,sha256=9JF40JYlW1y_AOqdD1pwYZJpJWZn63SC0K0VX8IA2JU,661
-x_transformers-1.37.9.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.37.9.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.37.9.dist-info/RECORD,,
+x_transformers-1.37.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.37.10.dist-info/METADATA,sha256=qA7YAj5ZeaesnGamBR-cPSR_0HSwquwPBptaUmi7c3c,662
+x_transformers-1.37.10.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+x_transformers-1.37.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.37.10.dist-info/RECORD,,