broccoli-ml 0.26.0__tar.gz → 0.28.0__tar.gz

{broccoli_ml-0.26.0 → broccoli_ml-0.28.0}/PKG-INFO

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

{broccoli_ml-0.26.0 → broccoli_ml-0.28.0}/broccoli/transformer.py

@@ -238,6 +238,7 @@ class FeedforwardBlock(nn.Module):
         linear_module=nn.Linear,
         pre_norm=True,
         normformer=False,
+        raw_input=False,
     ):
         super().__init__()
 
@@ -246,6 +247,11 @@ class FeedforwardBlock(nn.Module):
         else:
             self.activation = activation()
 
+        if raw_input:
+            self.memory_type = AnchoredLinear
+        else:
+            self.memory_type = nn.Linear
+
         self.dropout = nn.Dropout(dropout)
 
         self.max_features = (
@@ -260,7 +266,7 @@ class FeedforwardBlock(nn.Module):
                 linear_module(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                linear_module(ratio * output_features, output_features),
+                self.memory_type(ratio * output_features, output_features),
                 self.dropout,
             ]
         )
@@ -371,9 +377,11 @@ class TransformerBlock(nn.Module):
             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)
-            x = self.layer_norm_2(
-                torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
-            )
+
+        # Always post norm as eventually we reach the classification head!
+        x = self.layer_norm_2(
+            torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
+        )
 
         return x
 

{broccoli_ml-0.26.0 → broccoli_ml-0.28.0}/broccoli/vit.py

@@ -308,6 +308,7 @@ class ViTEncoder(nn.Module):
                 linear_module=linear_module,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
+                raw_input=not cnn,
             )
         else:
             self.initial_ff = nn.Identity()

{broccoli_ml-0.26.0 → broccoli_ml-0.28.0}/pyproject.toml

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