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

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/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/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.17.0.dist-info → broccoli_ml-0.18.0.dist-info}/METADATA

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

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

@@ -1,5 +1,5 @@
 broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
-broccoli/activation.py,sha256=jmKSNcq3VfZdVm8Ed65iiUB0ZfqmP_7lmEGkAWSIMdQ,2519
+broccoli/activation.py,sha256=AYTH6BNgXSBN8fnOMec_x94gISFj_r9y1IXIvhur7I4,2425
 broccoli/assets/2025_resnet_imagenet_1k_pretrained_state_dict.pkl,sha256=RZpPupWxFaVfgZrK-gBgfW1hj78oMEGhVWTbjRB3qMo,46835797
 broccoli/assets/cifar100_eigenvectors_size_2.pt,sha256=DjXDOXMeuMpIqNuGhX9z-OWYVqZwIMScSXZApRr9JjU,2501
 broccoli/assets/cifar100_eigenvectors_size_3.pt,sha256=gL6k0xtXYiYP6ZSvEiMBdJ7kIkT0AngTpDJHFQqwgxA,7173
@@ -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=E2JK5mQwJf75e23-JGcDoT7QxQf89DJReUo2et1LhRY,1716
+broccoli/tensor.py,sha256=txJiMLXTgWKTDd7Dh2In-ajiQLRUuuPezO0bKnEU3e8,1719
 broccoli/transformer.py,sha256=ZJpZIR5d4W_N1uO2AtRVVYzRJSjq1T6lj2mPT6raLg4,15923
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=LkMLqjLXG5SnNMeWLQ1LFfuG-RvQ5YY7BLZ-fNrVcPY,16607
-broccoli_ml-0.17.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.17.0.dist-info/METADATA,sha256=VPOAFsxDXoFtOQQkZoeWmUxZ9aRzyu1I9LX7N3WrXUA,1257
-broccoli_ml-0.17.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.17.0.dist-info/RECORD,,
+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,,