broccoli-ml 0.24.1__tar.gz → 0.24.3__tar.gz

{broccoli_ml-0.24.1 → broccoli_ml-0.24.3}/PKG-INFO

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

{broccoli_ml-0.24.1 → broccoli_ml-0.24.3}/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_ml-0.24.1 → broccoli_ml-0.24.3}/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_ml-0.24.1 → broccoli_ml-0.24.3}/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.1 → broccoli_ml-0.24.3}/pyproject.toml

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