broccoli-ml 0.37.0__tar.gz → 0.39.0__tar.gz

{broccoli_ml-0.37.0 → broccoli_ml-0.39.0}/PKG-INFO

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

{broccoli_ml-0.37.0 → broccoli_ml-0.39.0}/broccoli/linear.py

@@ -34,7 +34,8 @@ class SpectralNormLinear(nn.Module):
 
     def reset_parameters(self) -> None:
         weights = torch.empty(self.out_features, self.in_features)
-        nn.init.kaiming_uniform_(weights, a=math.sqrt(5))
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
         if self.use_bias:
             fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
             bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
@@ -77,7 +78,8 @@ class AnchoredLinear(nn.Module):
 
     def reset_parameters(self) -> None:
         weights = torch.empty(self.out_features, self.in_features)
-        nn.init.kaiming_uniform_(weights, a=math.sqrt(5))
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
         if self.use_bias:
             fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
             bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
@@ -120,7 +122,8 @@ class WeightNormedLinear(nn.Module):
 
     def reset_parameters(self) -> None:
         weights = torch.empty(self.out_features, self.in_features)
-        nn.init.kaiming_uniform_(weights, a=math.sqrt(5))
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
         if self.use_bias:
             fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
             bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
@@ -131,8 +134,7 @@ class WeightNormedLinear(nn.Module):
         return F.linear(x, self.weights(), self.bias)
 
     def __repr__(self) -> str:
-        # Optional: A nice representation for printing the module.
         return (
-            f"AnchoredLinear(in_features={self.in_features},"
+            f"WeightNormedLinear(in_features={self.in_features},"
             f"out_features={self.out_features}, bias={self.use_bias})"
         )

{broccoli_ml-0.37.0 → broccoli_ml-0.39.0}/broccoli/transformer.py

@@ -303,6 +303,8 @@ class TransformerBlock(nn.Module):
         mlp_ratio=4,
         activation: nn.Module = nn.ReLU,
         activation_kwargs: Optional[dict] = None,
+        ff_linear_module_up=None,
+        ff_linear_module_down=None,
         mlp_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
@@ -356,8 +358,16 @@ class TransformerBlock(nn.Module):
             activation=activation,
             activation_kwargs=activation_kwargs,
             dropout=mlp_dropout,
-            linear_module_up=linear_module,
-            linear_module_down=linear_module,
+            linear_module_up=(
+                ff_linear_module_up
+                if ff_linear_module_up is not None
+                else linear_module
+            ),
+            linear_module_down=(
+                ff_linear_module_down
+                if ff_linear_module_down is not None
+                else linear_module
+            ),
             pre_norm=False,  # Handled outside the block
             normformer=normformer,
             post_norm=False,  # Handled outside the block
@@ -389,51 +399,6 @@ class TransformerBlock(nn.Module):
 
         return x
 
-    #     if not self.training:
-    #         identity_probability = 0.0
-    #     else:
-    #         identity_probability = self.identity_probability
-
-    #     if random.random() < identity_probability:
-    #         return x
-    #     else:
-    #         ...
-
-    #     # perform the identity operation for some rows in the batch
-    #     dist = torch.distributions.Binomial(x.size(0), identity_probability)
-    #     identity_count = int(dist.sample().item())
-
-    #     shuffle_indices = torch.randperm(x.size(0), device=x.device)
-    #     unshuffle_indices = torch.argsort(shuffle_indices)
-    #     shuffled = x[shuffle_indices, :, :]
-    #     norm_shuffled = self.layer_norm_1(shuffled)
-    #     identity_x = shuffled[:identity_count, :, :]
-    #     process_x = shuffled[identity_count:, :, :]
-    #     residual = process_x
-
-    #     if self.pre_norm:
-    #         process_x = norm_shuffled[identity_count:, :, :]
-
-    #     process_x = residual + self.attn(process_x, process_x, process_x)
-    #     residual = process_x
-
-    #     shuffled = torch.cat([identity_x, process_x])
-    #     norm_shuffled = self.layer_norm_2(shuffled)
-
-    #     if self.pre_norm:
-    #         residual = process_x # residual NOT normed
-    #         process_x = norm_shuffled[identity_count:, :, :]
-
-    #     if self.post_norm:
-    #         process_x = norm_shuffled[identity_count:, :, :]
-    #         residual = process_x # residual normed
-
-    #     process_x = residual + self.ff(process_x) # handles residual connection
-
-    #     x = torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
-
-    #     return x if not self.post_norm else self.layer_norm_3(x)
-
 
 class TransformerEncoder(nn.Module):
     """
@@ -452,6 +417,8 @@ class TransformerEncoder(nn.Module):
         mlp_ratio=4,
         activation: nn.Module = nn.ReLU,
         activation_kwargs: Optional[dict] = None,
+        ff_linear_module_up=None,
+        ff_linear_module_down=None,
         mlp_dropout=0.0,
         msa_dropout=0.0,
         stochastic_depth=0.0,
@@ -515,6 +482,8 @@ class TransformerEncoder(nn.Module):
                     mlp_ratio=mlp_ratio,
                     activation=activation,
                     activation_kwargs=activation_kwargs,
+                    ff_linear_module_up=ff_linear_module_up,
+                    ff_linear_module_down=ff_linear_module_down,
                     mlp_dropout=mlp_dropout,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],

{broccoli_ml-0.37.0 → broccoli_ml-0.39.0}/broccoli/vit.py

@@ -132,6 +132,8 @@ class ViTEncoder(nn.Module):
         transformer_return_bos_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
+        transformer_ff_linear_module_up=None,
+        transformer_ff_linear_module_down=None,
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
@@ -282,6 +284,8 @@ class ViTEncoder(nn.Module):
                 mlp_ratio=transformer_mlp_ratio,
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
+                ff_linear_module_up=transformer_ff_linear_module_up,
+                ff_linear_module_down=transformer_ff_linear_module_down,
                 mlp_dropout=transformer_mlp_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
@@ -305,14 +309,16 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
                 linear_module_up=(
+                    # First truthy assigned value
                     transformer_initial_ff_linear_module_up
-                    if transformer_initial_ff_linear_module_up is not None
-                    else linear_module
+                    or transformer_ff_linear_module_up
+                    or linear_module
                 ),
                 linear_module_down=(
+                    # First truthy assigned value
                     transformer_initial_ff_linear_module_down
-                    if transformer_initial_ff_linear_module_down is not None
-                    else linear_module
+                    or transformer_ff_linear_module_down
+                    or linear_module
                 ),
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
@@ -386,6 +392,8 @@ class ViT(nn.Module):
         transformer_return_bos_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
+        transformer_ff_linear_module_up=None,
+        transformer_ff_linear_module_down=None,
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
@@ -446,6 +454,8 @@ class ViT(nn.Module):
             transformer_return_bos_tokens=transformer_return_bos_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,
+            transformer_ff_linear_module_up=transformer_ff_linear_module_up,
+            transformer_ff_linear_module_down=transformer_ff_linear_module_down,
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,

{broccoli_ml-0.37.0 → broccoli_ml-0.39.0}/pyproject.toml

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