broccoli-ml 0.18.1__tar.gz → 0.20.0__tar.gz

{broccoli_ml-0.18.1 → broccoli_ml-0.20.0}/PKG-INFO

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

{broccoli_ml-0.18.1 → broccoli_ml-0.20.0}/broccoli/tensor.py

@@ -14,30 +14,36 @@ class SigmaReparamTensor(nn.Module):
 
         super().__init__()
 
-        self.tensor = nn.Parameter(init_tensor, requires_grad=True)
+        self.sigma_reparam_tensor = nn.Parameter(init_tensor, requires_grad=True)
 
         with torch.no_grad():
-            _, sigma, v_transpose = torch.linalg.svd(self.tensor, full_matrices=False)
+            _, sigma, v_transpose = torch.linalg.svd(
+                self.sigma_reparam_tensor, full_matrices=False
+            )
 
         self.register_buffer("approx_spectral_norm", sigma[:1])
         self.register_buffer("right_singular", v_transpose[0])
-        self.scale = nn.Parameter(
+        self.sigma_reparam_scale = nn.Parameter(
             self.approx_spectral_norm.clone().detach(), requires_grad=True
         )
 
     def power_iteration(self):
         with torch.no_grad():
-            approx_right_singular_transpose = self.tensor.mv(self.right_singular)
+            approx_right_singular_transpose = self.sigma_reparam_tensor.mv(
+                self.right_singular
+            )
             approx_right_singular_transpose = F.normalize(
                 approx_right_singular_transpose, dim=0
             )
-            updated_right_singular = self.tensor.T.mv(approx_right_singular_transpose)
+            updated_right_singular = self.sigma_reparam_tensor.T.mv(
+                approx_right_singular_transpose
+            )
             updated_right_singular = F.normalize(updated_right_singular, dim=0)
             self.right_singular.data.copy_(updated_right_singular)
             rayleigh_quotient = torch.einsum(
                 "m,mn,n->",
                 approx_right_singular_transpose,
-                self.tensor,
+                self.sigma_reparam_tensor,
                 updated_right_singular,
             )
             self.approx_spectral_norm.data.copy_(rayleigh_quotient)
@@ -45,4 +51,6 @@ class SigmaReparamTensor(nn.Module):
     def forward(self):
         if self.training:
             self.power_iteration()
-        return self.scale * (self.tensor / self.approx_spectral_norm)
+        return self.sigma_reparam_scale * (
+            self.sigma_reparam_tensor / self.approx_spectral_norm
+        )

{broccoli_ml-0.18.1 → broccoli_ml-0.20.0}/broccoli/transformer.py

@@ -255,10 +255,8 @@ class FeedforwardBlock(nn.Module):
 
         if regularise_values:
             self.memory_type = SpectralNormLinear
-            self.bias_memories = False
         else:
             self.memory_type = nn.Linear
-            self.bias_memories = True
 
         self.process = nn.Sequential(
             *[
@@ -266,9 +264,7 @@ class FeedforwardBlock(nn.Module):
                 linear_module(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(ratio * output_features),
-                self.memory_type(
-                    ratio * output_features, output_features, bias=self.bias_memories
-                ),
+                self.memory_type(ratio * output_features, output_features),
                 self.dropout,
             ]
         )

{broccoli_ml-0.18.1 → broccoli_ml-0.20.0}/broccoli/vit.py

@@ -30,7 +30,7 @@ class GetCLSToken(nn.Module):
 
 
 class SequencePool(nn.Module):
-    def __init__(self, d_model, linear_module):
+    def __init__(self, d_model, linear_module=nn.Linear):
         super().__init__()
         self.attention = nn.Sequential(
             *[
@@ -54,13 +54,6 @@ class ClassificationHead(nn.Module):
         super().__init__()
         self.d_model = d_model
         self.summarize = GetCLSToken()
-        self.process = nn.Sequential(
-            *[
-                linear_module(d_model, 1),
-                Rearrange("batch seq 1 -> batch seq"),
-                nn.Softmax(dim=-1),
-            ]
-        )
         self.projection = nn.Linear(d_model, n_classes)
         if batch_norm:
             self.batch_norm = nn.BatchNorm1d(n_classes, affine=False)
@@ -87,7 +80,7 @@ class SequencePoolClassificationHead(ClassificationHead):
     """
 
     def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
-        super().__init__(d_model, linear_module, out_dim, batch_norm=True)
+        super().__init__(d_model, linear_module, out_dim, batch_norm=batch_norm)
         self.summarize = SequencePool(d_model, linear_module)
 
 

{broccoli_ml-0.18.1 → broccoli_ml-0.20.0}/pyproject.toml

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