broccoli-ml 0.24.0__py3-none-any.whl → 0.24.2__py3-none-any.whl

broccoli/activation.py

@@ -8,16 +8,18 @@ class ReLU(nn.Module):
     A ReLU activation function with optional clamp and leakiness.
     """
 
-    def __init__(self, clamp=True, leaky=True, leaky_slope=0.01, clamp_max=6.0) -> None:
+    def __init__(
+        self, clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0
+    ) -> None:
         super().__init__()
         self.clamp = clamp
         self.leaky = leaky
-        self.leaky_slope = leaky_slope
+        self.negative_slope = negative_slope
         self.clamp_max = clamp_max
 
     def forward(self, x):
         if self.leaky:
-            relu = F.leaky_relu(x, leaky_slope=self.leaky_slope)
+            relu = F.leaky_relu(x, negative_slope=self.negative_slope)
         else:
             relu = F.relu(x)
         if self.clamp:
@@ -69,17 +71,17 @@ class SquaredReLU(nn.Module):
     """
 
     def __init__(
-        self, clamp=True, leaky=True, leaky_slope: float = 0.01, clamp_max=6
+        self, clamp=True, leaky=True, negative_slope: float = 0.01, clamp_max=6
     ) -> None:
         super().__init__()
         self.clamp = clamp
         self.leaky = leaky
-        self.leaky_slope = leaky_slope
+        self.negative_slope = negative_slope
         self.clamp_max = clamp_max
 
     def forward(self, x):
         if self.leaky:
-            relu = F.leaky_relu(x, leaky_slope=self.leaky_slope)
+            relu = F.leaky_relu(x, negative_slope=self.negative_slope)
         else:
             relu = F.relu(x)
         relu_squared = relu**2
@@ -102,12 +104,12 @@ class XGLU(nn.Module):
         return self.activation(gate) * value
 
 
-def SquaredReGLU(clamp=True, leaky=True, leaky_slope=0.01, clamp_max=6.0) -> XGLU:
+def SquaredReGLU(clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0) -> XGLU:
     """
     Factory function that creates a GLU with a SquaredReLU activation.
     """
     activation_module = SquaredReLU(
-        clamp=clamp, leaky=leaky, leaky_slope=leaky_slope, clamp_max=clamp_max
+        clamp=clamp, leaky=leaky, negative_slope=negative_slope, clamp_max=clamp_max
     )
     return XGLU(activation_module)
 

broccoli/tensor.py

@@ -76,21 +76,27 @@ class AnchoredReparamTensor(nn.Module):
         assert init_tensor.ndim == 2, "Input tensor must be a 2D matrix."
         super().__init__()
 
-        self.weight = nn.Parameter(init_tensor.clone(), requires_grad=True)
+        # Use the gradboard convention of calling something nondecay_* if we should
+        # exclude it from weight decay
+        self.nondecay_weight = nn.Parameter(init_tensor.clone(), requires_grad=True)
 
         # At initialization, compute the dominant right-singular vector (v_0)
         # and store it in a non-trainable buffer.
         with torch.no_grad():
-            _, _, v_transpose = torch.linalg.svd(self.weight, full_matrices=False)
+            _, _, v_transpose = torch.linalg.svd(
+                self.nondecay_weight, full_matrices=False
+            )
             # v_transpose[0] is the first row of V^T, which is the first right-singular vector.
             self.register_buffer("anchor_vector", v_transpose[0])
 
-        initial_norm = torch.linalg.vector_norm(self.weight.mv(self.anchor_vector))
+        initial_norm = torch.linalg.vector_norm(
+            self.nondecay_weight.mv(self.anchor_vector)
+        )
         self.scale = nn.Parameter(initial_norm.clone().detach(), requires_grad=True)
 
     def forward(self) -> torch.Tensor:
         # Calculate the L2 norm of the matrix-vector product W @ v_0
-        norm = torch.linalg.vector_norm(self.weight.mv(self.anchor_vector))
+        norm = torch.linalg.vector_norm(self.nondecay_weight.mv(self.anchor_vector))
 
         # Return the reparameterized tensor.
-        return self.scale * (self.weight / (norm + 1e-6))
+        return self.scale * (self.nondecay_weight / (norm + 1e-6))

broccoli/transformer.py

@@ -236,7 +236,7 @@ class FeedforwardBlock(nn.Module):
         activation_kwargs=None,
         dropout=0.0,
         linear_module=nn.Linear,
-        reparam=False,
+        raw_input=False,
     ):
         super().__init__()
 
@@ -253,8 +253,9 @@ class FeedforwardBlock(nn.Module):
             else ratio * output_features
         )
 
-        if reparam:
+        if raw_input:
             self.memory_type = AnchoredLinear
+
         else:
             self.memory_type = linear_module
 
@@ -263,7 +264,7 @@ class FeedforwardBlock(nn.Module):
                 nn.LayerNorm(input_features),
                 linear_module(input_features, self.max_features),
                 self.activation,
-                # nn.LayerNorm(ratio * output_features),
+                nn.LayerNorm(ratio * output_features) if raw_input else nn.Identity(),
                 self.memory_type(ratio * output_features, output_features),
                 self.dropout,
             ]

broccoli/vit.py

@@ -302,7 +302,7 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
                 linear_module=linear_module,
-                reparam=not cnn,
+                raw_input=not cnn,
             )
         else:
             self.initial_ff = nn.Identity()

{broccoli_ml-0.24.0.dist-info → broccoli_ml-0.24.2.dist-info}/METADATA

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

{broccoli_ml-0.24.0.dist-info → broccoli_ml-0.24.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
-broccoli/activation.py,sha256=aBEEiWnvfyr8xKE47YWrncJqPEu9ymhd8uYd0s-_CeM,3385
+broccoli/activation.py,sha256=-Jf30C6iGqWCorC9HEGn2oduWwjeaCAxGLUUYIy1zX8,3438
 broccoli/assets/2025_resnet_imagenet_1k_pretrained_state_dict.pkl,sha256=RZpPupWxFaVfgZrK-gBgfW1hj78oMEGhVWTbjRB3qMo,46835797
 broccoli/assets/cifar100_eigenvectors_size_2.pt,sha256=DjXDOXMeuMpIqNuGhX9z-OWYVqZwIMScSXZApRr9JjU,2501
 broccoli/assets/cifar100_eigenvectors_size_3.pt,sha256=gL6k0xtXYiYP6ZSvEiMBdJ7kIkT0AngTpDJHFQqwgxA,7173
@@ -7,11 +7,11 @@ broccoli/cnn.py,sha256=jeRyKIAMWu1E3iyI14MGgSZuZivPMh12iqkqW9ilNjo,17785
 broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=4bxVDsO8E1d5-RZ23u160ZntazrT7Vt4AYTdAdCQU-w,3300
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
-broccoli/tensor.py,sha256=_YJP9tSFRkoKrR7cfnROSpWqfMyJLjgPmtFxEWRwgz8,3606
-broccoli/transformer.py,sha256=L1bVQZLUbtFtOy30yPVkjnqyELGhQoHJ_lFP_WPfYUA,16073
+broccoli/tensor.py,sha256=BVRS9-IWYtKXbXlerULQqitaUT8TPNgi7QTT5G_SmY4,3826
+broccoli/transformer.py,sha256=67y7FDfWAI-cIu8w5gRuG0MX24z1zflB218b1NKvk9A,16108
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=qGCx4cnpAkPpVHFrz6bFHdnPJXPaCxtTxKlI9YQJZWg,15649
-broccoli_ml-0.24.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.24.0.dist-info/METADATA,sha256=XwjOu6kQ0tKTUN0uc8H60YL4UvTPEyr2RYrfd4YGiyc,1257
-broccoli_ml-0.24.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.24.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=pCjoJo6d1EhWmzI45sdP8t0olu6FlCGwMkGXcJJXXlo,15651
+broccoli_ml-0.24.2.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.24.2.dist-info/METADATA,sha256=591Hk6OMRrWRzjyPS6To4t2uxOFa468XaPB9MlLSzZQ,1257
+broccoli_ml-0.24.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.24.2.dist-info/RECORD,,