x-transformers 2.6.5__py3-none-any.whl → 2.6.6__py3-none-any.whl

x_transformers/attend.py

@@ -176,7 +176,7 @@ class Attend(Module):
         softclamp_logits = False,
         logit_softclamp_value = 50.,
         add_zero_kv = False,
-        head_learned_sinks = 0,
+        head_learned_sink = False,
         selective = False,
         hard = False,
         cope = None,
@@ -257,10 +257,10 @@ class Attend(Module):
 
         # learned sink concatted pre-softmax, working solution from gpt-oss
 
-        self.has_head_learned_sinks = head_learned_sinks > 0
-        assert not (self.has_head_learned_sinks and flash), f'not supported for flash attention yet'
+        assert not (head_learned_sink and flash), f'not supported for flash attention yet'
 
-        self.head_attn_sinks = Parameter(torch.zeros(heads, head_learned_sinks)) if self.has_head_learned_sinks else None
+        self.head_learned_sink = head_learned_sink
+        self.head_attn_sink = Parameter(torch.zeros(heads)) if head_learned_sink else None
 
         # soft clamp attention logit value
 
@@ -517,10 +517,9 @@ class Attend(Module):
         if self.selective:
             sim = selective_attn(sim)
 
-        if self.has_head_learned_sinks:
+        if self.head_learned_sink:
             # add learned attention sink
-            num_sinks = self.head_attn_sinks.shape[-1]
-            attn_sink = repeat(self.head_attn_sinks, 'h sinks -> b h i sinks', b = sim.shape[0], i = sim.shape[2])
+            attn_sink = repeat(self.head_attn_sink, 'h -> b h i 1', b = sim.shape[0], i = sim.shape[2])
             sim = cat((attn_sink, sim), dim = -1)
 
         pre_softmax_attn = sim
@@ -531,9 +530,9 @@ class Attend(Module):
 
         post_softmax_attn = attn
 
-        if self.has_head_learned_sinks:
+        if self.head_learned_sink:
             # remove attention sink
-            attn = attn[..., num_sinks:]
+            attn = attn[..., 1:]
 
         attn = self.attn_dropout(attn)
 

x_transformers/x_transformers.py

@@ -1319,7 +1319,7 @@ class Attention(Module):
         value_dim_head = None,
         dim_out = None,
         add_zero_kv = False,         # same as add_zero_attn in pytorch
-        head_learned_sinks = 0,
+        head_learned_sink = False,
         rotate_num_heads = None,
         data_dependent_alibi = False,
         data_dependent_alibi_per_row = False,
@@ -1516,7 +1516,7 @@ class Attention(Module):
             selective = selective,
             custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,
-            head_learned_sinks = head_learned_sinks,
+            head_learned_sink = head_learned_sink,
             flash = flash,
             softclamp_logits = softclamp_logits,
             logit_softclamp_value = logit_softclamp_value,

{x_transformers-2.6.5.dist-info → x_transformers-2.6.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.6.5
+Version: 2.6.6
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.6.5.dist-info → x_transformers-2.6.6.dist-info}/RECORD

@@ -1,5 +1,5 @@
 x_transformers/__init__.py,sha256=aVuhUU0572TJHW88BVc4yA2tla0Zb8l3NH7W4RZ1AEs,1005
-x_transformers/attend.py,sha256=DX_qrDkz98Db0eNapbzciJbVp5dsWIFWdpv2LUfebJs,18223
+x_transformers/attend.py,sha256=JJv6ypJbZIFmH1LQ49hFg6hD0Wf9Z7Im1AP2ekm9hVI,18091
 x_transformers/autoregressive_wrapper.py,sha256=BsGO9xfVYkvynqbU1__tu_S_cxl7gss0YwnkhIa2baY,18401
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=hpb1sSbt3k2LNzzjrjSd8F5xOIbKj7IluV9MBEAFLkw,13031
@@ -9,10 +9,10 @@ x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=F_ZR9jysYmkbqKvsZmzXqOP3VznVeivXVOstAwKIdPU,122185
+x_transformers/x_transformers.py,sha256=vjRMEMA12Js94YwLVeZksYMEoRgK6CSKT6TJViMPp7U,122186
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.6.5.dist-info/METADATA,sha256=yMl0MlBbo7D9dOu_cBQz38iJQ3a6F8PlaPCo5RQXrSA,90445
-x_transformers-2.6.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.6.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.6.5.dist-info/RECORD,,
+x_transformers-2.6.6.dist-info/METADATA,sha256=95CKrJ98X7R0hpb5D8GHSfi372UtxXDSeDaO2qB0Lrs,90445
+x_transformers-2.6.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.6.6.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.6.6.dist-info/RECORD,,