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

x_transformers/attend.py

@@ -36,6 +36,9 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
+def at_most_one_of(*bools):
+    return sum([*map(int, bools)]) <= 1
+
 def compact(arr):
     return [*filter(exists, arr)]
 
@@ -100,6 +103,7 @@ class Attend(Module):
         scale = None,
         qk_norm = False,
         l2_distance = False,
+        sigmoid = False,
         flash = False,
         softclamp_logits = False,
         logit_softclamp_value = 50.,
@@ -116,10 +120,25 @@ class Attend(Module):
         super().__init__()
         self.scale = scale
 
+        # causal related
+
         self.causal = causal
         self.create_causal_mask = onnx_create_causal_mask if onnxable else create_causal_mask
 
-        self.attn_fn = partial(F.softmax, dtype = torch.float32) if not qk_norm else F.softmax
+        # attention type
+
+        assert not (flash and sigmoid), 'sigmoid attention not available for flash'
+        assert at_most_one_of(sigmoid, l2_distance)
+
+        self.sigmoid = sigmoid
+
+        if 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:
+            self.attn_fn = F.sigmoid
+
+        # dropouts
 
         self.dropout = dropout
         self.attn_dropout = nn.Dropout(dropout)
@@ -211,12 +230,12 @@ class Attend(Module):
 
         if self.l2_distance:
             k_norm_sq = k.norm(dim = -1, keepdim = True) ** 2
-            k = F.pad(k, (0, 1), value = 1.)
-            k = torch.cat((k, -k_norm_sq), dim = -1)
+            k = F.pad(k, (0, 1), value = -1.)
+            k = torch.cat((k, k_norm_sq), dim = -1)
 
             q_norm_sq = q.norm(dim = -1, keepdim = True) ** 2
-            q = torch.cat((2 * q, -q_norm_sq), dim = -1)
-            q = F.pad(q, (0, 1), value = 1.)
+            q = torch.cat((2 * q, q_norm_sq), dim = -1)
+            q = F.pad(q, (0, 1), value = -1.)
 
         # handle scale - by default they scale by dim_head ** -0.5, but need to take care if using cosine sim attention
 
@@ -410,7 +429,7 @@ class Attend(Module):
         if self.sigsoftmax:
             sim = sim + sim.sigmoid().log()
 
-        attn = self.attn_fn(sim, dim = -1)
+        attn = self.attn_fn(sim)
         attn = attn.type(dtype)
 
         post_softmax_attn = attn

x_transformers/x_transformers.py

@@ -924,6 +924,7 @@ class Attention(Module):
         qk_norm_scale = 10,
         qk_norm_dim_scale = False,
         l2_distance = False,
+        sigmoid = False,
         one_kv_head = False,
         kv_heads = None,
         shared_kv = False,
@@ -1039,6 +1040,7 @@ class Attention(Module):
             qk_norm = qk_norm,
             scale = qk_norm_scale if qk_norm else self.scale,
             l2_distance = l2_distance,
+            sigmoid = sigmoid,
             add_zero_kv = add_zero_kv,
             flash = flash,
             softclamp_logits = softclamp_logits,

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

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

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

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=BOnMjgV5O5DyAM_bUz1rI6n1j_eLXu8GIEljT-MMnWU,13434
+x_transformers/attend.py,sha256=EtaTN6ahgRlFLkwfHA31RNL_bQyAHwhNBpGU1NIHJ-c,13894
 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=-2fj6QcDSfMI5lJA_fzOW2mdzdS1C1LD6jMBtGQY48E,83752
+x_transformers/x_transformers.py,sha256=Y77_ZPWSKTJ-oYk4bHjhwMEkgoMaq_LyxcmCkvOPZ9g,83808
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.37.7.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.37.7.dist-info/METADATA,sha256=svdldk1hpiBN4xxKkOaQuMEXAxvaj6fT7Ri9NXwZJCU,661
-x_transformers-1.37.7.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.37.7.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.37.7.dist-info/RECORD,,
+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,,