broccoli-ml 0.16.0__tar.gz → 0.18.0__tar.gz

{broccoli_ml-0.16.0 → broccoli_ml-0.18.0}/PKG-INFO

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

{broccoli_ml-0.16.0 → broccoli_ml-0.18.0}/broccoli/activation.py

@@ -12,11 +12,11 @@ class SwiGLU(nn.Module):
     Halves the incoming parameter count, which should be scaled up before input.
     """
 
-    def __init__(self, linear_module: nn.Module = nn.Linear) -> None:
+    def __init__(self) -> None:
         super().__init__()
         # Learnable parameter is called "swiglu beta" so that it is easy to find
         #   and exclude from weight decay
-        self.swiglu_beta = nn.Parameter(torch.tensor([0.0]))
+        self.swiglu_beta = nn.Parameter(torch.tensor([1.0]))
 
     def forward(self, x):
         gate, value = rearrange(x, "... (split c) -> split ... c", split=2)
@@ -38,11 +38,13 @@ class SquaredReLU(nn.Module):
         self.leaky = leaky
 
     def forward(self, x):
-        relu_squared = F.relu(x) ** 2
+        if self.leaky:
+            relu = F.leaky_relu(x)
+        else:
+            relu = F.relu(x)
+        relu_squared = relu**2
         if self.clamp:
             relu_squared = torch.clamp(relu_squared, max=6)
-        if self.leaky:
-            relu_squared = relu_squared + 0.1 * x
         return relu_squared
 
 
@@ -57,11 +59,12 @@ class ReLU(nn.Module):
         self.leaky = leaky
 
     def forward(self, x):
-        relu = F.relu(x)
+        if self.leaky:
+            relu = F.leaky_relu(x)
+        else:
+            relu = F.relu(x)
         if self.clamp:
             relu = torch.clamp(relu, max=6)
-        if self.leaky:
-            relu = relu + 0.1 * x
         return relu
 
 
@@ -70,16 +73,13 @@ class GELU(nn.Module):
     A ReLU activation function with optional clamp and leakiness.
     """
 
-    def __init__(self, clamp=True, leaky=True) -> None:
+    def __init__(self, clamp=True) -> None:
         super().__init__()
         self.clamp = clamp
-        self.leaky = leaky
         self.gelu = nn.GELU()
 
     def forward(self, x):
         gelu = self.gelu(x)
         if self.clamp:
             gelu = torch.clamp(gelu, max=6)
-        if self.leaky:
-            gelu = gelu + 0.1 * x
         return gelu

{broccoli_ml-0.16.0 → broccoli_ml-0.18.0}/broccoli/tensor.py

@@ -32,7 +32,7 @@ class SigmaReparamTensor(nn.Module):
                 approx_right_singular_transpose, dim=0
             )
             updated_right_singular = self.tensor.T.mv(approx_right_singular_transpose)
-            updated_right_singular = F.normalize(self.right_singular, dim=0)
+            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->",

{broccoli_ml-0.16.0 → broccoli_ml-0.18.0}/broccoli/transformer.py

@@ -258,7 +258,7 @@ class FeedforwardBlock(nn.Module):
                 linear_module(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(ratio * output_features),
-                linear_module(ratio * output_features, output_features, bias=False),
+                linear_module(ratio * output_features, output_features),
                 self.dropout,
             ]
         )

{broccoli_ml-0.16.0 → broccoli_ml-0.18.0}/broccoli/vit.py

@@ -176,6 +176,9 @@ class ViTEncoder(nn.Module):
             )
 
         if cnn:
+            # This block rhymes:
+            if cnn_activation.__name__.endswith("GLU"):
+                cnn_out_channels *= 2
             cnn_output_size = calculate_output_spatial_size(
                 input_size,
                 kernel_size=cnn_kernel_size,
@@ -206,9 +209,6 @@ class ViTEncoder(nn.Module):
                     batchnormxd(cnn_activation_out_channels),
                 ]
             )
-            # This block rhymes:
-            if cnn and cnn_activation.__name__.endswith("GLU"):
-                cnn_out_channels *= 2
         else:
             self.cnn = nn.Identity()
             self.activate_and_dropout = nn.Identity()

{broccoli_ml-0.16.0 → broccoli_ml-0.18.0}/pyproject.toml

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