x-transformers 2.7.5__tar.gz → 2.7.6__tar.gz

{x_transformers-2.7.5 → x_transformers-2.7.6}/PKG-INFO

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

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

{x_transformers-2.7.5 → x_transformers-2.7.6}/x_transformers/x_transformers.py

@@ -2469,12 +2469,15 @@ class AttentionLayers(Module):
     ):
         # 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):
+        layer_and_layer_types = (self.layers, self.layer_types)
 
-            if layer_type not in ('a', 'c'):
-                continue
+        attn_layers = [layer for (_, layer, _), layer_type in zip(self.layers, self.layer_types) if layer_type in ('a', 'c')]
+        attn_intermeds = intermediates.attn_intermediates
+
+        assert len(attn_layers) == len(attn_intermeds)
 
-            layer.qk_clip_(attn_inter, tau = tau)
+        for attn_layer, attn_inter in zip(attn_layers, attn_intermeds):
+            attn_layer.qk_clip_(attn_inter, tau = tau)
 
     def forward(
         self,