broccoli-ml 13.0.5__tar.gz → 14.0.0__tar.gz

{broccoli_ml-13.0.5 → broccoli_ml-14.0.0}/PKG-INFO

@@ -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 → broccoli_ml-14.0.0}/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 → broccoli_ml-14.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "13.0.5"
+version = "14.0.0"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}