x-transformers 1.23.1__py3-none-any.whl → 1.23.2__py3-none-any.whl

x_transformers/x_transformers.py

@@ -650,6 +650,7 @@ class Attention(nn.Module):
         num_mem_kv = 0,
         dropout = 0.,
         on_attn = False,
+        gate_value_heads = False,
         gate_values = False,
         zero_init_output = False,
         max_attend_past = None,
@@ -703,7 +704,14 @@ class Attention(nn.Module):
         if gate_values:
             self.to_v_gate = nn.Linear(dim, out_dim)
             nn.init.constant_(self.to_v_gate.weight, 0)
-            nn.init.constant_(self.to_v_gate.bias, 1)
+            nn.init.constant_(self.to_v_gate.bias, 10)
+
+        # add per head gating of the output values, from 'Attend to nothing' paper
+        self.to_v_head_gate = None
+        if gate_value_heads:
+            self.to_v_head_gate = nn.Linear(dim, heads)
+            nn.init.constant_(self.to_v_head_gate.weight, 0)
+            nn.init.constant_(self.to_v_head_gate.bias, 10)
 
         # cosine sim attention
         self.qk_norm = qk_norm
@@ -905,6 +913,12 @@ class Attention(nn.Module):
         if head_scale:
             out = out * self.head_scale_params
 
+        # per head gating, from https://arxiv.org/abs/2306.12929
+
+        if exists(self.to_v_head_gate):
+            head_gate = self.to_v_head_gate(x)
+            out = out * rearrange(head_gate, 'b n h -> b h n 1').sigmoid()
+
         # merge heads
 
         out = rearrange(out, 'b h n d -> b n (h d)')

{x_transformers-1.23.1.dist-info → x_transformers-1.23.2.dist-info}/METADATA

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

{x_transformers-1.23.1.dist-info → x_transformers-1.23.2.dist-info}/RECORD

@@ -3,10 +3,10 @@ x_transformers/attend.py,sha256=hZcz_iijzbEqbXp2_BPEVL-1LoHXmYaHE6e6Oy-7hFE,1126
 x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
 x_transformers/continuous_autoregressive_wrapper.py,sha256=pTiDqu6JRUlnQJQp_xHATYHy0lgSd6ERLqyiFO3pC-4,1575
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=ikf99q_g1_v_wObZed972s2hHrbiDpAq_qGJDmNcVZc,60573
+x_transformers/x_transformers.py,sha256=KQ9mU_jE27whl6yQI67grF0S8Xhd3GndnM6Yd0-q-lw,61162
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers-1.23.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.23.1.dist-info/METADATA,sha256=91Gu0qU9ztioZ2_oVeOrLeVkP--n6ngneDxEOMUHJe8,661
-x_transformers-1.23.1.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.23.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.23.1.dist-info/RECORD,,
+x_transformers-1.23.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.23.2.dist-info/METADATA,sha256=8h0sbx8-4yNTOJuAZLbe5HQ16hsmZI1M_mT-rMIIMJc,661
+x_transformers-1.23.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.23.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.23.2.dist-info/RECORD,,