broccoli-ml 12.0.0__tar.gz → 12.1.0__tar.gz

{broccoli_ml-12.0.0 → broccoli_ml-12.1.0}/PKG-INFO

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

{broccoli_ml-12.0.0 → broccoli_ml-12.1.0}/broccoli/transformer.py

@@ -592,33 +592,52 @@ class TransformerBlock(nn.Module):
         return self.attn._kv_distance
 
     def forward(self, x):
-
-        if self.pre_norm:
-            process_x = self.pre_attention_norm(x)
-        else:
-            process_x = x
-
-        processed = self.drop_path(self.attn(process_x, process_x, process_x))
-
-        if self.normformer:
-            processed = self.normformer_norm(processed)
-
-        x = x + processed
-
-        if self.post_norm:
-            x = self.post_attention_norm(x)
-
         if self.pre_norm:
-            process_x = self.pre_mlp_norm(x)
-        else:
-            process_x = x
-
-        x = x + self.drop_path(self.ff(process_x))
-
-        if self.post_norm:
-            x = self.post_mlp_norm(x)
-
-        return x
+            x = self.layer_norm_1(x)
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = self.layer_norm_2(x)
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
+            if self.post_norm:  # i.e. in addition! Pre and post.
+                x = self.layer_norm_3(x)
+        elif self.post_norm:  # i.e. only, not prenorm, just post
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = self.layer_norm_1(x)
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
+            x = self.layer_norm_2(x)
+        else:  # Not pre or post norm. Stand well back.
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
+
+        # if self.pre_norm:
+        #     process_x = self.pre_attention_norm(x)
+        # else:
+        #     process_x = x
+
+        # processed = self.drop_path(self.attn(process_x, process_x, process_x))
+
+        # if self.normformer:
+        #     processed = self.normformer_norm(processed)
+
+        # if self.residual_path:
+        #     x = x + processed
+
+        # if self.post_norm:
+        #     x = self.post_attention_norm(x)
+
+        # if self.pre_norm:
+        #     process_x = self.pre_mlp_norm(x)
+        # else:
+        #     process_x = x
+
+        # processed = self.drop_path(self.ff(process_x))
+
+        # if self.residual_path:
+        #     x = x + processed
+
+        # if self.post_norm:
+        #     x = self.post_mlp_norm(x)
+
+        # return x
 
     def attention_logits(self, x):
         """

{broccoli_ml-12.0.0 → broccoli_ml-12.1.0}/pyproject.toml

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