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

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

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

{broccoli_ml-0.18.1 → broccoli_ml-0.19.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.19.0}/pyproject.toml

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