broccoli-ml 12.0.0__py3-none-any.whl → 12.1.0__py3-none-any.whl

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.dist-info → broccoli_ml-12.1.0.dist-info}/METADATA

@@ -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.dist-info → broccoli_ml-12.1.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=dKiqPZEBLK5e6hqEhz99Z68RCJZNAPkTwQQQfjFDhko,27611
+broccoli/transformer.py,sha256=cPXHGlfkvJAXgr_ONyESQ0RAEzn7yqZEoyZ3l4N3EX8,28571
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
 broccoli/vit.py,sha256=DvVpayMIcUhH7Xg6CiYyeedUuuMHrjsGxEdXfnTGa_Q,22428
-broccoli_ml-12.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-12.0.0.dist-info/METADATA,sha256=xt43usBeTgqLUqnHqIsPf8xWN3Bkd8ckRP1GE69YID0,1369
-broccoli_ml-12.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-12.0.0.dist-info/RECORD,,
+broccoli_ml-12.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-12.1.0.dist-info/METADATA,sha256=2j4nt5aHVcPiW_YF2Sx4Yu_l3MKe2fa0S9v3aQ1PZE0,1369
+broccoli_ml-12.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-12.1.0.dist-info/RECORD,,