x-transformers 1.39.0__py3-none-any.whl → 1.39.3__py3-none-any.whl

x_transformers/attend.py

@@ -107,6 +107,15 @@ def qk_l2_dist_squared(q, k):
     l2_dist_squared = torch.cdist(q, k) ** 2
     return unpack_one(l2_dist_squared, packed_shape, '* i j')
 
+# gumbel softmax
+
+def gumbel_softmax(t, temperature = 1.):
+    one_hot_indices = t.argmax(dim = -1, keepdim = True)
+    one_hot = torch.zeros_like(t).scatter(-1, one_hot_indices, 1.)
+
+    t = (t / temperature).softmax(dim = -1)
+    return one_hot + t - t.detach()
+
 # functions for creating causal mask
 # need a special one for onnx cpu (no support for .triu)
 
@@ -128,8 +137,9 @@ class Attend(Module):
         dropout = 0.,
         causal = False,
         heads = None,
-        pre_talking_heads = True,
-        post_talking_heads = True,
+        pre_talking_heads = False,
+        post_talking_heads = False,
+        pre_scale_post_talking_heads = False,
         sparse_topk = None,
         scale = None,
         qk_norm = False,
@@ -141,6 +151,7 @@ class Attend(Module):
         logit_softclamp_value = 50.,
         add_zero_kv = False,
         selective = False,
+        hard = False,
         sigsoftmax = False,
         cope = None,
         onnxable = False,
@@ -161,17 +172,18 @@ class Attend(Module):
         # attention type
 
         assert not (flash and sigmoid), 'sigmoid attention not available for flash'
-        assert at_most_one_of(sigmoid, l2_distance)
-
-        self.sigmoid = sigmoid
+        assert not (flash and hard), 'hard attention not available for flash'
+        assert at_most_one_of(sigmoid, hard, l2_distance)
 
         if exists(custom_attn_fn):
             self.attn_fn = custom_attn_fn
-        elif not sigmoid:
+        elif sigmoid:
+            self.attn_fn = F.sigmoid
+        elif hard:
+            self.attn_fn = gumbel_softmax
+        else:
             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
 
@@ -180,10 +192,11 @@ class Attend(Module):
 
         # talking heads
 
-        assert not (flash and (pre_talking_heads or post_talking_heads)), 'talking heads not compatible with flash attention'
+        assert not (flash and (pre_talking_heads or post_talking_heads or pre_scale_post_talking_heads)), 'talking heads not compatible with flash attention'
 
         self.pre_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if pre_talking_heads else None
         self.post_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if post_talking_heads else None
+        self.pre_scale_post_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if pre_scale_post_talking_heads else None
 
         if exists(self.pre_softmax_talking_heads):
             nn.init.dirac_(self.pre_softmax_talking_heads.weight)
@@ -191,6 +204,10 @@ class Attend(Module):
         if exists(self.post_softmax_talking_heads):
             nn.init.dirac_(self.post_softmax_talking_heads.weight)
 
+        if exists(self.pre_scale_post_talking_heads):
+            # an improvisation where heads are combined pre-softmax attention, then used to scale post-softmax attention
+            nn.init.dirac_(self.pre_scale_post_talking_heads.weight)
+
         # selective attention
 
         assert not (flash and selective), 'selective attention cannot work on flash attention'
@@ -436,6 +453,9 @@ class Attend(Module):
 
         qk_similarities = sim.clone()
 
+        if exists(self.pre_scale_post_talking_heads):
+            pre_to_post_scale = self.pre_scale_post_talking_heads(sim)
+
         if exists(self.pre_softmax_talking_heads):
             sim = sim + self.pre_softmax_talking_heads(sim)
 
@@ -478,6 +498,7 @@ class Attend(Module):
             sim = sim + sim.sigmoid().log()
 
         attn = self.attn_fn(sim)
+
         attn = attn.type(dtype)
 
         post_softmax_attn = attn
@@ -487,6 +508,9 @@ class Attend(Module):
         if exists(self.post_softmax_talking_heads):
             attn = self.post_softmax_talking_heads(attn)
 
+        if exists(self.pre_scale_post_talking_heads):
+            attn = attn * pre_to_post_scale
+
         out = einsum(f'b h i j, {kv_einsum_eq} -> b h i d', attn, v)
 
         intermediates = Intermediates(

x_transformers/x_transformers.py

@@ -909,6 +909,7 @@ class Attention(Module):
         flash = False,
         pre_talking_heads = False,
         post_talking_heads = False,
+        pre_scale_post_talking_heads = False,
         head_scale = False,
         sparse_topk = None,
         num_mem_kv = 0,
@@ -918,6 +919,7 @@ class Attention(Module):
         swiglu_values = False,
         gate_values = False,
         zero_init_output = False,
+        hard = False,
         sigsoftmax = False,
         max_attend_past = None,
         qk_norm = False,
@@ -1039,8 +1041,10 @@ class Attention(Module):
             causal = causal,
             pre_talking_heads = pre_talking_heads,
             post_talking_heads = post_talking_heads,
+            pre_scale_post_talking_heads = pre_scale_post_talking_heads,
             dropout = dropout,
             sparse_topk = sparse_topk,
+            hard = hard,
             qk_norm = qk_norm,
             scale = qk_norm_scale if qk_norm else self.scale,
             l2_distance = l2_distance,
@@ -1610,6 +1614,7 @@ class AttentionLayers(Module):
                 layer = post_branch_fn(layer)
 
             residual_fn = GRUGating if gate_residual else Residual
+
             residual = residual_fn(dim, scale_residual = scale_residual, scale_residual_constant = scale_residual_constant)
 
             # handle unet skip connection

{x_transformers-1.39.0.dist-info → x_transformers-1.39.3.dist-info}/METADATA

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

{x_transformers-1.39.0.dist-info → x_transformers-1.39.3.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=sTeX7DmUt6I5FhHtgcTDOIvmD1CvJ1PmVjZ_-lYO-QA,15596
+x_transformers/attend.py,sha256=K4DNcWAgOHzNtOTIJA15VA3VQ2KMyv-PX8oO1R0Z5Rw,16670
 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=nfq-EOLx5HWf8tXlmVuDbkhqNfFnfqRMCEkALK3SFkA,84341
+x_transformers/x_transformers.py,sha256=3JrVqwYVrd5UVf2esdunTcer7QL72H7VF4mL3UsCWOI,84508
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.39.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.39.0.dist-info/METADATA,sha256=-ppCMMH6ZTsmwaMJB9q4b4Yvd-nU8v95l5SdzTY17OU,661
-x_transformers-1.39.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.39.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.39.0.dist-info/RECORD,,
+x_transformers-1.39.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.39.3.dist-info/METADATA,sha256=aqCockVfuZfLM3D9ZdlgN_HfUPdhZm4UWyKg2HLkuUo,661
+x_transformers-1.39.3.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+x_transformers-1.39.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.39.3.dist-info/RECORD,,

{x_transformers-1.39.0.dist-info → x_transformers-1.39.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (75.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any