broccoli-ml 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

broccoli/transformer.py

@@ -299,9 +299,12 @@ class TransformerBlock(nn.Module):
     ):
         super().__init__()
 
+        self.pre_norm = pre_norm
+
         self.identity_probability = identity_probability
 
-        self.layer_norm = nn.LayerNorm(d_model)
+        self.layer_norm_1 = nn.LayerNorm(d_model)
+        self.layer_norm_2 = nn.LayerNorm(d_model)
 
         if position_embedding_type == "relative":
             max_freq = int(max(source_size) / 2)  # Suggested by Gemini!
@@ -358,12 +361,21 @@ class TransformerBlock(nn.Module):
         identity_x = shuffled[:identity_count, :, :]
         process_x = shuffled[identity_count:, :, :]
 
-        norm_process_x = self.layer_norm(process_x)
-        process_x = process_x + self.attn(
-            norm_process_x, norm_process_x, norm_process_x
+        if self.pre_norm:
+            norm_process_x = self.layer_norm_1(process_x)
+            process_x = process_x + self.attn(
+                norm_process_x, norm_process_x, norm_process_x
+            )
+            process_x = process_x + self.ff(process_x)
+        else:  # post-norm
+            process_x = process_x + self.attn(process_x, process_x, process_x)
+            norm_process_x = self.layer_norm_1(process_x)
+            process_x = process_x + self.ff(process_x)
+
+        # Always post norm as eventually we reach the classification head!
+        x = self.layer_norm_2(
+            torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
         )
-        process_x = process_x + self.ff(process_x)
-        x = torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
 
         return x
 

broccoli/vit.py

@@ -306,7 +306,8 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
                 linear_module=linear_module,
-                raw_input=not cnn,
+                pre_norm=transformer_pre_norm,
+                normformer=transformer_normformer,
             )
         else:
             self.initial_ff = nn.Identity()

{broccoli_ml-0.25.0.dist-info → broccoli_ml-0.27.0.dist-info}/METADATA

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

{broccoli_ml-0.25.0.dist-info → broccoli_ml-0.27.0.dist-info}/RECORD

@@ -8,10 +8,10 @@ broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=7NkNvhtxzWAUoBJOuiPUIcr853HhI1cS71d8DwdMkJ0,4826
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=zhSOo9W24FEgN7U35wy3ZIJHnw3u4cepJO5heCw6vwU,4590
-broccoli/transformer.py,sha256=BBnbmGwmvTmpBL4LDykUTeOHXGzIKNYdNIm6X9LEcGA,16278
+broccoli/transformer.py,sha256=barmcq4Y5X6iM6STFt2t52XlXhuZphQ98wX6UvUGMFU,16748
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=GNbwld8NtnDR2D-bBDrhuMEUWwBbH2KVq29nwmYdJto,16009
-broccoli_ml-0.25.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.25.0.dist-info/METADATA,sha256=qkm7knCzbtTK3Cdc6bT2l4tTlcmm2oIJ61n8AmZMLMM,1257
-broccoli_ml-0.25.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.25.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=nuXX2JoKoBTtF1tAH-11mL2R5ISgMHYsBbbgvcluV1s,16072
+broccoli_ml-0.27.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.27.0.dist-info/METADATA,sha256=dzd2nejaPCG976rKlMnDTi5RzpWAeh87Yf2LsPb2Tzw,1257
+broccoli_ml-0.27.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.27.0.dist-info/RECORD,,