broccoli-ml 0.18.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

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/transformer.py

@@ -236,6 +236,7 @@ class FeedforwardBlock(nn.Module):
         activation_kwargs=None,
         dropout=0.0,
         linear_module=nn.Linear,
+        regularise_values=True,
     ):
         super().__init__()
 
@@ -252,13 +253,22 @@ class FeedforwardBlock(nn.Module):
             else ratio * output_features
         )
 
+        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(
             *[
                 nn.LayerNorm(input_features),
                 linear_module(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(ratio * output_features),
-                linear_module(ratio * output_features, output_features),
+                self.memory_type(
+                    ratio * output_features, output_features, bias=self.bias_memories
+                ),
                 self.dropout,
             ]
         )

{broccoli_ml-0.18.0.dist-info → broccoli_ml-0.19.0.dist-info}/METADATA

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

{broccoli_ml-0.18.0.dist-info → broccoli_ml-0.19.0.dist-info}/RECORD

@@ -7,11 +7,11 @@ broccoli/cnn.py,sha256=jeRyKIAMWu1E3iyI14MGgSZuZivPMh12iqkqW9ilNjo,17785
 broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=g8YrxNl6g_WcHrWVmbaBHJU5hv6daFS0r4TxAoPJ9UE,3012
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
-broccoli/tensor.py,sha256=txJiMLXTgWKTDd7Dh2In-ajiQLRUuuPezO0bKnEU3e8,1719
-broccoli/transformer.py,sha256=ZJpZIR5d4W_N1uO2AtRVVYzRJSjq1T6lj2mPT6raLg4,15923
+broccoli/tensor.py,sha256=MUvXtwD2f1sPTBym4FB0x_ZfsJUBNLgULUlN8btV8GI,1943
+broccoli/transformer.py,sha256=2hoRg9N0g2HTfrUU6cTXl4msHO0orWiHS9xJTUMHej4,16234
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
 broccoli/vit.py,sha256=--eEAxBzCCVGCIvaZWHsonbVXGe04wDEz0Q3V4YZXoI,16599
-broccoli_ml-0.18.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.18.0.dist-info/METADATA,sha256=1CcIDoXMMzX88L6sqjEgtrkXJV4HSyQhRymaGBTfQkY,1257
-broccoli_ml-0.18.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.18.0.dist-info/RECORD,,
+broccoli_ml-0.19.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.19.0.dist-info/METADATA,sha256=aSUin_G9VrKRhl8gBqGirK1yWJyMo6krnWzA0KWF1tQ,1257
+broccoli_ml-0.19.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.19.0.dist-info/RECORD,,