broccoli-ml 0.34.0__py3-none-any.whl → 0.35.0__py3-none-any.whl

broccoli/transformer.py

@@ -235,7 +235,8 @@ class FeedforwardBlock(nn.Module):
         activation=nn.ReLU,
         activation_kwargs=None,
         dropout=0.0,
-        linear_module=nn.Linear,
+        linear_module_up=nn.Linear,
+        linear_module_down=nn.Linear,
         pre_norm=True,
         normformer=False,
         post_norm=True,
@@ -265,10 +266,10 @@ class FeedforwardBlock(nn.Module):
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
-                linear_module(input_features, self.max_features),
+                linear_module_up(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                linear_module(ratio * output_features, output_features),
+                linear_module_down(ratio * output_features, output_features),
                 self.dropout,
             ]
         )
@@ -346,7 +347,7 @@ class TransformerBlock(nn.Module):
             bos_tokens=bos_tokens,
         )
 
-        # Submodules for the feedforward process
+        # Submodule for the feedforward process
         self.ff = FeedforwardBlock(
             d_model,
             mlp_ratio,
@@ -354,7 +355,8 @@ class TransformerBlock(nn.Module):
             activation=activation,
             activation_kwargs=activation_kwargs,
             dropout=mlp_dropout,
-            linear_module=linear_module,
+            linear_module_up=linear_module,
+            linear_module_down=linear_module,
             pre_norm=pre_norm,
             normformer=normformer,
             post_norm=post_norm,

broccoli/vit.py

@@ -118,7 +118,8 @@ class ViTEncoder(nn.Module):
         pooling_padding=1,
         transformer_feedforward_first=True,
         transformer_initial_ff_residual_path=True,
-        transformer_initial_ff_linear_module=None,
+        transformer_initial_ff_linear_module_up=None,
+        transformer_initial_ff_linear_module_down=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -296,6 +297,7 @@ class ViTEncoder(nn.Module):
                 return_bos_tokens=transformer_return_bos_tokens,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
+                post_norm=transformer_post_norm,
             )
         else:
             self.transformer = nn.Identity()
@@ -308,13 +310,19 @@ class ViTEncoder(nn.Module):
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
-                linear_module=(
-                    transformer_initial_ff_linear_module
-                    if transformer_initial_ff_linear_module is not None
+                linear_module_up=(
+                    transformer_initial_ff_linear_module_up
+                    if transformer_initial_ff_linear_module_up is not None
+                    else linear_module
+                ),
+                linear_module_down=(
+                    transformer_initial_ff_linear_module_down
+                    if transformer_initial_ff_linear_module_down is not None
                     else linear_module
                 ),
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
+                post_norm=transformer_post_norm,
                 residual_path=transformer_initial_ff_residual_path,
             )
         else:
@@ -370,7 +378,8 @@ class ViT(nn.Module):
         pooling_padding=1,
         transformer_feedforward_first=True,
         transformer_initial_ff_residual_path=True,
-        transformer_initial_ff_linear_module=None,
+        transformer_initial_ff_linear_module_up=None,
+        transformer_initial_ff_linear_module_down=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -429,7 +438,8 @@ class ViT(nn.Module):
             pooling_padding=pooling_padding,
             transformer_feedforward_first=transformer_feedforward_first,
             transformer_initial_ff_residual_path=transformer_initial_ff_residual_path,
-            transformer_initial_ff_linear_module=transformer_initial_ff_linear_module,
+            transformer_initial_ff_linear_module_up=transformer_initial_ff_linear_module_up,
+            transformer_initial_ff_linear_module_down=transformer_initial_ff_linear_module_down,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
             transformer_post_norm=transformer_post_norm,

{broccoli_ml-0.34.0.dist-info → broccoli_ml-0.35.0.dist-info}/METADATA

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

{broccoli_ml-0.34.0.dist-info → broccoli_ml-0.35.0.dist-info}/RECORD

@@ -8,10 +8,10 @@ broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=8Y9vD85ZEgNZsIQgO3uRQ3lOQR-JjwvabY8liCrfNCk,4831
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=ks2TRCdS10k2XvxEieh2sj_LzjTNRuiO6gekKFTtziI,4533
-broccoli/transformer.py,sha256=aF452hNruRRBZCqBqb0FM-C2tHH0sB5Fan937Plh7Cc,17106
+broccoli/transformer.py,sha256=t0gsADJC9UOlwjm7tDKdy0pAZ8l3clTcCnes86zvH-k,17203
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=HEdG9PPra3295Gv7jwzJbtiaRhJMp8QKFwv_dUzF67o,16827
-broccoli_ml-0.34.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.34.0.dist-info/METADATA,sha256=Fzkgbyz9r3VyBLP90AD_Vfv9DyDORiEoKsi0ipJDJ8w,1257
-broccoli_ml-0.34.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.34.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=c-ZRHiLDOoQDJO9OJ51zD9HqaluG33flIwTXQQfms-g,17389
+broccoli_ml-0.35.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.35.0.dist-info/METADATA,sha256=v0JSpcubSGwxA5dFPbDwz2r2oGZWSeqYND1Mu8WOiJY,1257
+broccoli_ml-0.35.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.35.0.dist-info/RECORD,,