broccoli-ml 0.10.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

broccoli/linear.py

@@ -1,9 +1,56 @@
 # UNDER CONSTRUCTION
 
+import math
 import torch
 from torch import nn
 from torch.nn import functional as F
 
+from .tensor import SigmaReparamTensor
+
+
+class SpectralNormLinear(nn.Module):
+    """
+    ...
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        self.weight_init = nn.Parameter(torch.empty(out_features, in_features))
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight_init, a=math.sqrt(5))
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = SigmaReparamTensor(self.weight_init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        # Optional: A nice representation for printing the module.
+        return (
+            f"SpectralNormFeedForward(in_features={self.in_features}",
+            f"out_features={self.out_features}, bias={self.use_bias})",
+        )
+
 
 class RandomLinear(nn.Linear):
     """ """

broccoli/transformer.py

@@ -10,6 +10,7 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from .rope import RotaryEmbedding, apply_rotary_emb
+from .linear import SpectralNormLinear
 
 
 class MHAttention(nn.Module):
@@ -235,6 +236,7 @@ class FeedforwardLayer(nn.Module):
         activation_kwargs=None,
         dropout=0.0,
         linear_module=nn.Linear,
+        norm_memory=False,
     ):
         super().__init__()
 
@@ -245,19 +247,28 @@ class FeedforwardLayer(nn.Module):
 
         self.dropout = nn.Dropout(dropout)
 
+        self.max_features = (
+            2 * ratio * output_features
+            if activation.__name__.endswith("GLU")
+            else ratio * output_features
+        )
+
+        if norm_memory:
+            self.memory_type = SpectralNormLinear
+            self.bias_memories = False
+        else:
+            self.memory_type = linear_module
+            self.bias_memories = True
+
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features),
-                linear_module(
-                    input_features,
-                    (
-                        2 * ratio * input_features
-                        if activation.__name__.endswith("GLU")
-                        else ratio * input_features
-                    ),
-                ),
+                linear_module(input_features, self.max_features),
                 self.activation,
-                linear_module(ratio * input_features, output_features),
+                nn.LayerNorm(self.max_features),
+                self.memory_type(
+                    ratio * output_features, output_features, bias=self.bias_memories
+                ),
                 self.dropout,
             ]
         )

broccoli/vit.py

@@ -82,6 +82,7 @@ class ViTEncoder(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         intermediate_feedforward_layer=True,
+        norm_intermediate_ff_memory=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -263,6 +264,7 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
                 linear_module=linear_module,
+                norm_memory=norm_intermediate_ff_memory,
             )
         elif pooling_out_channels < transformer_embedding_size:
             self.intermediate_feedforward_layer = nn.Identity()
@@ -326,6 +328,7 @@ class CCT(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         intermediate_feedforward_layer=True,
+        norm_intermediate_ff_memory=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,

{broccoli_ml-0.10.0.dist-info → broccoli_ml-0.12.0.dist-info}/METADATA

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

{broccoli_ml-0.10.0.dist-info → broccoli_ml-0.12.0.dist-info}/RECORD

@@ -5,13 +5,13 @@ broccoli/assets/cifar100_eigenvectors_size_2.pt,sha256=DjXDOXMeuMpIqNuGhX9z-OWYV
 broccoli/assets/cifar100_eigenvectors_size_3.pt,sha256=gL6k0xtXYiYP6ZSvEiMBdJ7kIkT0AngTpDJHFQqwgxA,7173
 broccoli/cnn.py,sha256=jeRyKIAMWu1E3iyI14MGgSZuZivPMh12iqkqW9ilNjo,17785
 broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
-broccoli/linear.py,sha256=0XYCi3ckTEKwAgBOMUSJP2HsnrroOH8eyrhRdpANG2w,1298
+broccoli/linear.py,sha256=9ZwqC6kkgkr0uPoEjdi_Uq1QFHb4wCXzuU1r2pDreXM,2910
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=E2JK5mQwJf75e23-JGcDoT7QxQf89DJReUo2et1LhRY,1716
-broccoli/transformer.py,sha256=RSZpbHs_K4ts5os6lWxcGDI3p0zreRwQNnk6mV8HJnk,15930
+broccoli/transformer.py,sha256=niooSyrG9kZPk60IUPa-ZevEGaUQ8MI6AUxuOInAoqc,16265
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=uXqMIvAVY4PuA-Fv1YxU7L3_74fR19GtNu1caQeMr6k,15185
-broccoli_ml-0.10.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.10.0.dist-info/METADATA,sha256=zRP8Hn-Q0uMDvNhInEuwcm-sVwG5zE3vsWZBSuVqDmE,1257
-broccoli_ml-0.10.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.10.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=J1-59oROlpbIEZ-grrXMhuluM3cLoJIAvixnIRSmgOs,15326
+broccoli_ml-0.12.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.12.0.dist-info/METADATA,sha256=ApiCIlFOB8rrdnUwLXgYv_0GPsQx85L3Uq2rUh3JhjQ,1257
+broccoli_ml-0.12.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.12.0.dist-info/RECORD,,