x-transformers 2.7.3__tar.gz → 2.7.4__tar.gz

{x_transformers-2.7.3 → x_transformers-2.7.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.7.3
+Version: 2.7.4
 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 → x_transformers-2.7.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.7.3"
+version = "2.7.4"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.7.3 → x_transformers-2.7.4}/tests/test_x_transformers.py

@@ -1314,3 +1314,14 @@ def test_simple_mdlm(
 
     loss = nar(seq)
     loss.loss.backward()
+
+def test_qk_clip():
+    from x_transformers import Attention
+
+    x = torch.randn(1, 1024, 512)
+
+    attn = Attention(dim = 512, dim_out = 384)
+
+    out, intermediates = attn(x, return_intermediates = True)
+
+    attn.qk_clip_(intermediates, tau = 100)

{x_transformers-2.7.3 → x_transformers-2.7.4}/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,