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

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,

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

@@ -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.dist-info → x_transformers-2.7.6.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=xaGBkYCy6CqL0q9icWmL_WzCeU6ZztEYEkMtN71L2z4,124576
+x_transformers/x_transformers.py,sha256=odnCZAKZKrQLXmpaWhiPVB5elGjt8kerDbO3-yeC-60,124764
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-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,,
+x_transformers-2.7.6.dist-info/METADATA,sha256=n-AKJXX2Ko3XlehMOv5MojPrFaHdRi4lRkvcGAFOXR4,93739
+x_transformers-2.7.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.7.6.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.7.6.dist-info/RECORD,,