broccoli-ml 13.0.6__tar.gz → 14.0.1__tar.gz

{broccoli_ml-13.0.6 → broccoli_ml-14.0.1}/PKG-INFO

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

{broccoli_ml-13.0.6 → broccoli_ml-14.0.1}/broccoli/transformer.py

@@ -612,15 +612,18 @@ class EncoderBlock(nn.Module):
         else:
             process_x = x
 
-        processed = self.drop_path(self.attn(process_x, process_x, process_x))
+        processed = self.attn(process_x, process_x, process_x)
 
         if self.normformer:
             processed = self.normformer_norm(processed)
 
+        processed = self.drop_path(processed)
+
         x = self.alpha * x + processed
 
         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 +641,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.6 → broccoli_ml-14.0.1}/pyproject.toml

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