x-transformers 2.7.3__py3-none-any.whl → 2.7.5__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1637,10 +1637,12 @@ class Attention(Module):
         q_weight = self.to_q.weight
         k_weight = self.to_k.weight
 
-        q_dim, k_dim, heads = q_weight.shape[0], k_weight.shape[0], qk_weight_scale.numel()
+        qk_dim, heads = q_weight.shape[0], qk_weight_scale.numel()
 
-        q_weight.mul_(repeat(qk_weight_scale, 'h -> (h expand)', expand = q_dim // heads))
-        k_weight.mul_(repeat(qk_weight_scale, 'h -> (h expand)', expand = k_dim // heads))
+        qk_weight_scale = repeat(qk_weight_scale, 'h -> (h expand)', expand = qk_dim // heads)
+
+        q_weight.mul_(qk_weight_scale)
+        k_weight.mul_(qk_weight_scale)
 
     def forward(
         self,
@@ -2460,6 +2462,20 @@ class AttentionLayers(Module):
 
         self.can_cache_kv = all([module.can_cache_kv for module in self.modules() if isinstance(module, Attention)])
 
+    def attn_qk_clip_(
+        self,
+        intermediates: LayerIntermediates,
+        tau = 100.
+    ):
+        # pairs up the attention intermediates with each attention module and does qk clip proposed by kimi team
+
+        for (_, layer, _), layer_type, attn_inter in zip(self.layers, self.layer_types, intermediates.attn_intermediates):
+
+            if layer_type not in ('a', 'c'):
+                continue
+
+            layer.qk_clip_(attn_inter, tau = tau)
+
     def forward(
         self,
         x,
@@ -3190,6 +3206,13 @@ class TransformerWrapper(Module):
             if not isinstance(self.pos_emb, always):
                 nn.init.normal_(self.pos_emb.emb.weight, std = 1e-5)
 
+    def attn_qk_clip_(
+        self,
+        intermediates: LayerIntermediates,
+        tau = 100.
+    ):
+        self.attn_layers.attn_qk_clip_(intermediates, tau = tau)
+
     def forward(
         self,
         x,

{x_transformers-2.7.3.dist-info → x_transformers-2.7.5.dist-info}/METADATA

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

{x_transformers-2.7.3.dist-info → x_transformers-2.7.5.dist-info}/RECORD

@@ -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=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=5lWGEmRqj2_XwxQW0CiXOikPdRUQpyICSecokJB19Mk,123971
+x_transformers/x_transformers.py,sha256=xaGBkYCy6CqL0q9icWmL_WzCeU6ZztEYEkMtN71L2z4,124576
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.7.3.dist-info/METADATA,sha256=XzaIZ5nLUiI_sGTz_uekMDZkOEqRsPxggpnm9SNcsfw,93739
-x_transformers-2.7.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.7.3.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.7.3.dist-info/RECORD,,
+x_transformers-2.7.5.dist-info/METADATA,sha256=m6f4PIgJFKKWlsGAydi_Bg5-7-0IRlor0pRY_zBh5s8,93739
+x_transformers-2.7.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.7.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.7.5.dist-info/RECORD,,