broccoli-ml 13.0.5__py3-none-any.whl → 14.0.0__py3-none-any.whl

broccoli/transformer.py

@@ -200,26 +200,26 @@ class MHAttention(nn.Module):
                 "`source_size` must be a tuple of 1, 2 or 3 integers"
             )
 
-        q = rearrange(q, "b t (h d) -> b t h d", h=self.n_heads)
-        k = rearrange(k, "b t (h d) -> b t h d", h=self.n_heads)
+        q = rearrange(q, "b t (h d) -> b h t d", h=self.n_heads)
+        k = rearrange(k, "b t (h d) -> b h t d", h=self.n_heads)
 
         q_util, q_img = (
-            q[:, : self.utility_tokens, :, :],
-            q[:, self.utility_tokens :, :, :],
+            q[:, :, : self.utility_tokens, :],
+            q[:, :, self.utility_tokens :, :],
         )
         k_util, k_img = (
-            k[:, : self.utility_tokens, :, :],
-            k[:, self.utility_tokens :, :, :],
+            k[:, :, : self.utility_tokens, :],
+            k[:, :, self.utility_tokens :, :],
         )
 
         q_img = rearrange(
             q_img,
-            f"b ({spatial_dimension_names}) h d -> b {spatial_dimension_names} h d",
+            f"b h ({spatial_dimension_names}) d -> b h {spatial_dimension_names} d",
             **spatial_dimension_values,
         )
         k_img = rearrange(
             k_img,
-            f"b ({spatial_dimension_names}) h d -> b {spatial_dimension_names} h d",
+            f"b h ({spatial_dimension_names}) d -> b h {spatial_dimension_names} d",
             **spatial_dimension_values,
         )
 
@@ -230,19 +230,19 @@ class MHAttention(nn.Module):
 
         q_img = rearrange(
             q_img,
-            f"b {spatial_dimension_names} h d -> b ({spatial_dimension_names}) h d",
+            f"b h {spatial_dimension_names} d -> b h ({spatial_dimension_names}) d",
         )
         k_img = rearrange(
             k_img,
-            f"b {spatial_dimension_names} h d -> b ({spatial_dimension_names}) h d",
+            f"b h {spatial_dimension_names} d -> b h ({spatial_dimension_names}) d",
         )
 
         # Re-combine the utility tokens and the RoPE-enhanced sequence tokens
-        q = torch.cat([q_util, q_img], dim=1)
-        k = torch.cat([k_util, k_img], dim=1)
+        q = torch.cat([q_util, q_img], dim=2)
+        k = torch.cat([k_util, k_img], dim=2)
 
-        q = rearrange(q, "b t h d -> b t (h d)")
-        k = rearrange(k, "b t h d -> b t (h d)")
+        q = rearrange(q, "b h t d -> b t (h d)")
+        k = rearrange(k, "b h t d -> b t (h d)")
 
         return q, k
 
@@ -621,6 +621,7 @@ class EncoderBlock(nn.Module):
 
         if self.post_norm:
             x = self.post_attention_norm(x)
+            process_x = x
         elif self.pre_norm:
             process_x = self.pre_mlp_norm(x)
         else:
@@ -638,15 +639,15 @@ class EncoderBlock(nn.Module):
     def attention_logits(self, x):
         """
         Give back the attention scores used in this layer.
+        Needs to match what the model actually sees during forward()
+        by applying the correct normalisations.
         """
-        # Fix: Use the correct attribute name 'pre_attention_norm'
         if self.pre_norm:
-            # We must normalize the input before measuring attention logits
-            # to match what the model actually sees during forward()
             x = self.pre_attention_norm(x)
-            return self.attn.attention_logits(x, x, x)
-        else:
-            return self.attn.attention_logits(x, x, x)
+        elif self.post_norm:
+            x = self.input_norm(x)
+
+        return self.attn.attention_logits(x, x, x)
 
     def reset_parameters(self):
         if self.pre_norm:

{broccoli_ml-13.0.5.dist-info → broccoli_ml-14.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 13.0.5
+Version: 14.0.0
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-13.0.5.dist-info → broccoli_ml-14.0.0.dist-info}/RECORD

@@ -4,10 +4,10 @@ broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
 broccoli/linear.py,sha256=W-3aNpBjd_0xRyzbCKkmg4H1qmslQOIQhB-WDDay2nM,13125
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=-rVhSl5yWDbEbjLysSpjiEy0h01E33TGdtgBnd-zRgA,27952
+broccoli/transformer.py,sha256=eSh_oyVhQad0rI0YPt1617OtH_EPn8XuNYj6BR6zRSM,27875
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
 broccoli/vit.py,sha256=SYAHGuGVZRa-zYCLK5siKdUauDgHtMYUmQ6W2tUIqKA,22588
-broccoli_ml-13.0.5.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-13.0.5.dist-info/METADATA,sha256=JM7ILFgCPGW7lkP_vWfet9KdNMz6xpRcNlChvzQA5xU,1369
-broccoli_ml-13.0.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-13.0.5.dist-info/RECORD,,
+broccoli_ml-14.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-14.0.0.dist-info/METADATA,sha256=XAj9Kyl3rn7tH3_Tm1CpnkfQeacS2w6ymjw4_AH73D4,1369
+broccoli_ml-14.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-14.0.0.dist-info/RECORD,,