broccoli-ml 0.1.0__tar.gz

broccoli_ml-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 nicholasbailey87
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

broccoli_ml-0.1.0/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.3
+Name: broccoli-ml
+Version: 0.1.0
+Summary: Some useful Pytorch models, circa 2025
+License: MIT
+Author: Nicholas Bailey
+Requires-Python: >=3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: einops (>=0.8.1,<0.9.0)
+Requires-Dist: numpy (>=2.3.1,<3.0.0)
+Requires-Dist: torch (>=2.7.1,<3.0.0)
+Description-Content-Type: text/markdown
+
+# broccoli
+Some useful Pytorch models, circa 2025
+

broccoli_ml-0.1.0/README.md

@@ -0,0 +1,2 @@
+# broccoli
+Some useful Pytorch models, circa 2025

broccoli_ml-0.1.0/broccoli/activation.py

@@ -0,0 +1,24 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from einops import rearrange
+
+
+class SwiGLU(nn.Module):
+    """
+    Implementation of (beta) SwiGLU, as introduced in "GLU Variants Improve Transformer"
+        (https://arxiv.org/abs/2002.05202v1) and used to great effect in LLaMa 2.0.
+
+    Halves the incoming parameter count, which should be scaled up before input.
+    """
+
+    def __init__(self, linear_module: nn.Module = nn.Linear) -> 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]))
+
+    def forward(self, x):
+        gate, value = rearrange(x, "... (split c) -> split ... c", split=2)
+        beta_swish = gate * F.sigmoid(self.swiglu_beta * gate)
+        return beta_swish * value

broccoli_ml-0.1.0/broccoli/cct.py

@@ -0,0 +1,249 @@
+from typing import Optional
+
+from transformer import TransformerEncoder
+from cnn import ConvLayer, ConcatPool, WhiteningConv
+from einops import einsum, rearrange
+from einops.layers.torch import Rearrange
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SequencePool(nn.Module):
+    """
+    As described in [Hasani et al. (2021) *''Escaping the Big Data Paradigm with
+        Compact Transformers''*](https://arxiv.org/abs/2104.05704). It can be viewed
+        as. ageneralisation of average pooling.
+    """
+
+    def __init__(self, d_model, linear_module, out_dim):
+        super().__init__()
+        self.d_model = d_model
+        self.attention = nn.Sequential(
+            *[
+                linear_module(d_model, 1),
+                Rearrange("batch seq 1 -> batch seq"),
+                nn.Softmax(dim=-1),
+            ]
+        )
+        self.projection = nn.Linear(d_model, out_dim)
+        self.norm = nn.BatchNorm1d(out_dim, affine=False)
+
+    def forward(self, x):
+        weights = self.attention(x)
+        weighted_embedding = einsum(
+            weights, x, "batch seq, batch seq d_model -> batch d_model"
+        )
+        projection = self.projection(weighted_embedding)
+        return self.norm(projection)
+
+
+class CCTEncoder(nn.Module):
+    """
+    Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
+        *''Escaping the Big Data Paradigm with Compact Transformers''*](
+        https://arxiv.org/abs/2104.05704). It's basically a convolutional neural
+        network leading into a transformer encoder. To make it like the full CCT
+        we would finish it of with a sequence pooling layer but we won't always
+        want to do that.
+    """
+
+    def __init__(
+        self,
+        image_size=32,
+        conv_kernel_size=3,  # Only 2 is supported for eigenvector initialisation
+        whitening=False,
+        pooling_type="maxpool",
+        pooling_kernel_size=3,
+        pooling_kernel_stride=2,
+        pooling_kernel_padding=1,
+        transformer_embedding_size=256,
+        transformer_layers=7,
+        transformer_heads=4,
+        transformer_mlp_ratio=2,
+        transformer_bos_tokens=1,
+        activation: nn.Module = nn.ReLU,
+        activation_kwargs: Optional[dict] = None,
+        mlp_dropout=0.0,
+        msa_dropout=0.1,
+        stochastic_depth=0.1,
+        linear_module=nn.Linear,
+        image_channels=3,
+    ):
+        # TODO: turn these into proper exceptions
+        if whitening and (conv_kernel_size != 2):
+            print("We currently only support whitening for kernel size 2!")
+            assert not (whitening and (conv_kernel_size != 2))
+        if pooling_type not in ["maxpool", "concat"]:
+            print("Pooling type must be maxpool or concat")
+            assert pooling_type in ["maxpool", "concat"]
+
+        if activation_kwargs is not None:
+            self.activation = activation(**activation_kwargs)
+        else:
+            self.activation = activation()
+
+        super().__init__()
+        self.image_size = image_size
+        output_size = (
+            (image_size + 2 * pooling_kernel_padding - pooling_kernel_size)
+            / pooling_kernel_stride
+        ) + 1
+
+        self.sequence_length = int(output_size) ** 2
+
+        if pooling_type == "maxpool":
+
+            conv_out_channels = transformer_embedding_size
+
+            self.pool = nn.Sequential(
+                *[
+                    Rearrange(
+                        "N C H W -> N H W C"
+                    ),  # rearrange in case we're using XGLU activation
+                    self.activation,
+                    Rearrange("N H W C -> N C H W"),
+                    nn.MaxPool2d(
+                        pooling_kernel_size,
+                        stride=pooling_kernel_stride,
+                        padding=pooling_kernel_padding,
+                    ),
+                    Rearrange("N C H W -> N (H W) C"),
+                ]
+            )
+        elif pooling_type == "concat":
+
+            conv_out_channels = int(
+                round(transformer_embedding_size / (pooling_kernel_size**2))
+            )
+
+            self.pool = ConcatPool(
+                conv_out_channels,
+                pooling_kernel_size,
+                pooling_kernel_stride,
+                pooling_kernel_padding,
+                transformer_embedding_size,
+                activation,
+                activation_kwargs,
+                linear_module,
+            )
+
+        if transformer_layers > 0:
+            self.transformer = TransformerEncoder(
+                self.sequence_length,
+                transformer_embedding_size,
+                transformer_layers,
+                transformer_heads,
+                mlp_ratio=transformer_mlp_ratio,
+                activation=activation,
+                activation_kwargs=activation_kwargs,
+                mlp_dropout=mlp_dropout,
+                msa_dropout=msa_dropout,
+                stochastic_depth=stochastic_depth,
+                causal=False,
+                linear_module=linear_module,
+                bos_tokens=transformer_bos_tokens,
+            )
+        else:
+            self.transformer = nn.Identity()
+
+        # This code block rhymes:
+        if activation.__name__.endswith("GLU"):
+            conv_out_channels *= 2
+
+        if whitening:
+            whitening_conv_out_channels = conv_kernel_size**2 * image_channels * 2
+            self.conv = nn.Sequential(
+                *[
+                    WhiteningConv(
+                        in_channels=image_channels,
+                        kernel_size=conv_kernel_size,
+                        linear_module=linear_module,
+                    ),
+                    Rearrange("N C H W -> N H W C"),
+                    nn.Linear(whitening_conv_out_channels, conv_out_channels),
+                    Rearrange("N H W C -> N C H W"),
+                ]
+            )
+        else:
+            self.conv = ConvLayer(
+                image_channels,
+                conv_out_channels,
+                kernel_size=conv_kernel_size,
+                stride=1,
+                padding="same",
+                linear_module=linear_module,
+            )
+
+        self.encoder = nn.Sequential(*[self.conv, self.pool, self.transformer])
+
+    def forward(self, x):
+        return self.encoder(x)
+
+
+class CCT(nn.Module):
+    """
+    Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
+        *''Escaping the Big Data Paradigm with Compact Transformers''*](
+        https://arxiv.org/abs/2104.05704). It's a convolutional neural network
+        leading into a transformer encoder, followed by a sequence pooling layer.
+    """
+
+    def __init__(
+        self,
+        image_size=32,
+        conv_kernel_size=3,  # Only 2 is supported for eigenvector initialisation
+        whitening=False,
+        pooling_type="maxpool",
+        pooling_kernel_size=3,
+        pooling_kernel_stride=2,
+        pooling_kernel_padding=1,
+        transformer_embedding_size=256,
+        transformer_layers=7,
+        transformer_heads=4,
+        transformer_mlp_ratio=2,
+        transformer_bos_tokens=1,
+        activation: nn.Module = nn.ReLU,
+        activation_kwargs: Optional[dict] = None,
+        mlp_dropout=0.0,  # The original paper got best performance from mlp_dropout=0.
+        msa_dropout=0.1,  # "" msa_dropout=0.1
+        stochastic_depth=0.1,  # "" stochastic_depth=0.1
+        image_classes=100,
+        linear_module=nn.Linear,
+        image_channels=3,
+    ):
+
+        assert not (whitening and (conv_kernel_size != 2))
+
+        super().__init__()
+
+        self.encoder = CCTEncoder(
+            image_size,
+            conv_kernel_size,
+            whitening,
+            pooling_type,
+            pooling_kernel_size,
+            pooling_kernel_stride,
+            pooling_kernel_padding,
+            transformer_embedding_size,
+            transformer_layers,
+            transformer_heads,
+            transformer_mlp_ratio,
+            transformer_bos_tokens,
+            activation=activation,
+            activation_kwargs=activation_kwargs,
+            mlp_dropout=mlp_dropout,
+            msa_dropout=msa_dropout,
+            stochastic_depth=stochastic_depth,
+            linear_module=linear_module,
+            image_channels=image_channels,
+        )
+        self.pool = SequencePool(
+            transformer_embedding_size, linear_module, image_classes
+        )
+
+    @property
+    def sequence_length(self):
+        return self.encoder.sequence_length
+
+    def forward(self, x):
+        return self.pool(self.encoder(x))

broccoli_ml-0.1.0/broccoli/cnn.py

@@ -0,0 +1,400 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import _pair
+from einops import rearrange, repeat
+import math
+from typing import Type, Union, Tuple, Optional, Literal
+
+from einops.layers.torch import Rearrange
+
+
+# Helper function to calculate padding for 'same' mode
+# Adapted from https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/conv.py
+def _calculate_same_padding(
+    input_size: Tuple[int, int],
+    kernel_size: Tuple[int, int],
+    stride: Tuple[int, int],
+    dilation: Tuple[int, int],
+) -> Tuple[int, int, int, int]:
+    """Calculates padding for 'same' output shape."""
+    ih, iw = input_size
+    kh, kw = kernel_size
+    sh, sw = stride
+    dh, dw = dilation
+
+    # Effective kernel size
+    eff_kh = (kh - 1) * dh + 1
+    eff_kw = (kw - 1) * dw + 1
+
+    # Calculate required total padding
+    out_h = (ih + sh - 1) // sh
+    out_w = (iw + sw - 1) // sw
+    pad_h = max((out_h - 1) * sh + eff_kh - ih, 0)
+    pad_w = max((out_w - 1) * sw + eff_kw - iw, 0)
+
+    # Distribute padding (similar to TensorFlow 'SAME' behavior)
+    pad_top = pad_h // 2
+    pad_bottom = pad_h - pad_top
+    pad_left = pad_w // 2
+    pad_right = pad_w - pad_left
+    return (pad_left, pad_right, pad_top, pad_bottom)
+
+
+# Custom Convolution Layer
+class ConvLayer(nn.Module):
+    """
+    A 2D Convolution layer implemented using torch.nn.Unfold and a custom linear layer.
+
+    This layer mimics the behavior of torch.nn.Conv2d but allows injecting
+    a different linear layer implementation for processing the unfolded patches.
+
+    Args:
+        in_channels (int): Number of channels in the input image.
+        out_channels (int): Number of channels produced by the convolution.
+        kernel_size (int or tuple): Size of the convolving kernel.
+        stride (int or tuple, optional): Stride of the convolution. Default: 1.
+        padding (int, tuple or str, optional): Padding added to all four sides
+            of the input. Can be an int, a tuple of two ints (padH, padW),
+            a tuple of four ints (padLeft, padRight, padTop, padBottom),
+            or the strings 'valid' (no padding) or 'same' (padding for same
+            output spatial dims as input). Default: 0 ('valid').
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1.
+        bias (bool, optional): If True, adds a learnable bias to the output.
+            The bias is handled by the underlying linear layer. Default: True.
+        linear (Type[nn.Module], optional): The class of the linear layer
+            to use for the kernel operation. Must accept (in_features, out_features, bias)
+            in its constructor. Defaults to torch.nn.Linear.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        stride: Union[int, Tuple[int, int]] = 1,
+        padding: Union[
+            int, Tuple[int, int], Tuple[int, int, int, int], Literal["valid", "same"]
+        ] = 0,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        bias: bool = True,
+        linear_module: Type[nn.Module] = nn.Linear,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        self.dilation = _pair(dilation)
+        self.bias = bias
+        self.linear_module = linear_module
+        self.padding_mode = (
+            padding  # Store the original padding mode ('same', 'valid', int, or tuple)
+        )
+
+        # Calculate the number of input features for the linear layer
+        # It's the number of channels times the kernel area
+        self.linear_in_features = (
+            in_channels * self.kernel_size[0] * self.kernel_size[1]
+        )
+
+        # Instantiate the linear layer (kernel)
+        self.kernel = self.linear_module(
+            self.linear_in_features, out_channels, bias=bias
+        )
+
+        # We will use F.pad for manual padding, so unfold padding is 0
+        self.unfold = nn.Unfold(
+            kernel_size=self.kernel_size,
+            dilation=self.dilation,
+            padding=0,  # Manual padding handled in forward
+            stride=self.stride,
+        )
+
+        # Determine numeric padding values for F.pad
+        if isinstance(padding, str):
+            if padding not in ["valid", "same"]:
+                raise ValueError("padding must be 'valid', 'same', an int, or a tuple")
+            # 'same' padding calculation depends on input size, defer to forward pass
+            # 'valid' padding means 0
+            self._padding_val = (
+                (0, 0, 0, 0) if padding == "valid" else None
+            )  # None indicates 'same'
+        elif isinstance(padding, int):
+            self._padding_val = (padding,) * 4
+        elif isinstance(padding, tuple) and len(padding) == 2:
+            # (padH, padW) -> (padW_left, padW_right, padH_top, padH_bottom)
+            self._padding_val = (padding[1], padding[1], padding[0], padding[0])
+        elif isinstance(padding, tuple) and len(padding) == 4:
+            # (padLeft, padRight, padTop, padBottom) - already in F.pad format
+            self._padding_val = padding
+        else:
+            raise TypeError(
+                "padding must be 'valid', 'same', an int, or a tuple of 2 or 4 ints"
+            )
+
+    def _calculate_output_shape(self, h_in: int, w_in: int) -> Tuple[int, int]:
+        """Calculates the output height and width."""
+        if self._padding_val is None:  # 'same' padding
+            # For 'same' padding, output size matches input size if stride is 1.
+            # If stride > 1, output size is ceil(input_size / stride)
+            # The _calculate_same_padding helper ensures this behavior.
+            oh = math.ceil(h_in / self.stride[0])
+            ow = math.ceil(w_in / self.stride[1])
+            return oh, ow
+        else:
+            # Use the standard formula with the calculated numeric padding
+            pad_h = self._padding_val[2] + self._padding_val[3]  # top + bottom
+            pad_w = self._padding_val[0] + self._padding_val[1]  # left + right
+            kh, kw = self.kernel_size
+            sh, sw = self.stride
+            dh, dw = self.dilation
+
+            eff_kh = (kh - 1) * dh + 1
+            eff_kw = (kw - 1) * dw + 1
+
+            oh = math.floor((h_in + pad_h - eff_kh) / sh + 1)
+            ow = math.floor((w_in + pad_w - eff_kw) / sw + 1)
+            return oh, ow
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Performs the forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (N, C_in, H_in, W_in).
+
+        Returns:
+            torch.Tensor: Output tensor of shape (N, C_out, H_out, W_out).
+        """
+        N, C, H, W = x.shape
+        if C != self.in_channels:
+            raise ValueError(
+                f"Input channels {C} does not match expected {self.in_channels}"
+            )
+
+        # 1. Calculate and Apply Padding
+        if self._padding_val is None:  # 'same' padding mode
+            pad_l, pad_r, pad_t, pad_b = _calculate_same_padding(
+                (H, W), self.kernel_size, self.stride, self.dilation
+            )
+            padded_x = F.pad(x, (pad_l, pad_r, pad_t, pad_b))
+            # Update H, W for output shape calculation after padding
+            # Note: _calculate_output_shape will correctly handle 'same' based on original H, W
+        elif self._padding_val != (0, 0, 0, 0):
+            padded_x = F.pad(x, self._padding_val)
+        else:  # No padding ('valid' or explicit 0)
+            padded_x = x
+
+        # 2. Unfold to extract patches
+        # Input: (N, C_in, H_pad, W_pad)
+        # Output: (N, C_in * K_h * K_w, L), where L is the number of patches (H_out * W_out)
+        patches = self.unfold(padded_x)
+        num_patches = patches.shape[-1]  # L
+
+        # 3. Reshape for the linear layer
+        # We want (N, L, C_in * K_h * K_w) to apply the linear layer patch-wise
+        # transpose switches the last two dimensions
+        patches_transposed = patches.transpose(1, 2)  # Shape: (N, L, C_in * K_h * K_w)
+
+        # 4. Apply the linear layer (kernel) to each patch
+        # Input: (N, L, linear_in_features)
+        # Output: (N, L, out_channels)
+        linear_output = self.kernel(patches_transposed)
+
+        # 5. Reshape back to image format
+        # We need (N, out_channels, L) first
+        output_transposed = linear_output.transpose(1, 2)  # Shape: (N, out_channels, L)
+
+        # Calculate output spatial dimensions
+        out_h, out_w = self._calculate_output_shape(H, W)  # Use original H, W
+
+        # Check if the number of patches matches the calculated output dimensions
+        if num_patches != out_h * out_w:
+            # This might happen with certain combinations of stride/padding/dilation/input size
+            # if the calculation logic has an issue. nn.Unfold is usually robust.
+            print(
+                f"Warning: Mismatch in calculated patches. "
+                f"Expected L={out_h * out_w}, got {num_patches}. "
+                f"Using unfolded L={num_patches} to determine output shape."
+            )
+            # Attempt recovery if possible, though might indicate upstream calculation error
+            # Find factors of num_patches close to expected out_h, out_w
+            # This part is tricky and might not always yield the desired shape.
+            # For simplicity, we'll rely on nn.Unfold's L and reshape.
+            # A more robust solution might re-calculate H_out, W_out based *only* on L.
+            # For now, let's stick to the reshape based on calculated out_h, out_w,
+            # assuming they match L. If they don't, the reshape will fail.
+            pass  # Proceed with calculated out_h, out_w
+
+        # Reshape using einops (or tensor.view)
+        # Input: (N, C_out, L) -> Output: (N, C_out, H_out, W_out)
+        output = rearrange(output_transposed, "n c (h w) -> n c h w", h=out_h, w=out_w)
+        # Alternative using view:
+        # output = output_transposed.view(N, self.out_channels, out_h, out_w)
+
+        return output
+
+    def extra_repr(self) -> str:
+        s = (
+            "{in_channels}, {out_channels}, kernel_size={kernel_size}"
+            ", stride={stride}"
+        )
+        if self.padding_mode != 0 and self.padding_mode != "valid":
+            s += ", padding={padding_mode}"
+        if self.dilation != (1,) * len(self.dilation):
+            s += ", dilation={dilation}"
+        # if self.groups != 1: # Not implemented
+        #     s += ', groups={groups}'
+        if self.bias is False:
+            s += ", bias=False"
+        if self.linear_module != nn.Linear:
+            s += f", linear={self.linear.__name__}"
+        return s.format(**self.__dict__)
+
+
+class WhiteningConv(ConvLayer):
+    def __init__(
+        self,
+        in_channels: int,
+        kernel_size: int,
+        eigenvectors: torch.Tensor,
+        bias: bool = True,
+        linear_module: Type[nn.Module] = nn.Linear,
+    ):
+        """
+        We end up using a concatenation of the eigenvector tensor with its negation,
+            as the tendency to use e.g. ReLU in neural networks means that useful
+            data may otherwise be lost (if one orientation of an eigenvector produces
+            a strong negative signal, this will be clipped to zero by ReLU, but a
+            strong positive signal from the negation of the eigenvector will be
+            preserved). Assuming a square kernel, out channels is thus
+
+            (kernel_size ** 2) * in_channels * 2
+
+            where the trailing "* 2" accounts for the doubling of the size of the
+            eigenvector tensor we're using by including the negative of each eigenvector
+            as well.
+        """
+        out_channels = kernel_size**2 * in_channels * 2
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            padding="same",
+            bias=bias,
+            linear_module=linear_module,
+        )
+        self.eigenvectors = torch.cat([eigenvectors, -eigenvectors], dim=0)
+        # bias updates if `bias`=True but weight doesn't,
+        #   per Jordan (2024) https://arxiv.org/abs/2404.00498
+        #   but weight is set to `requires_grad = False`:
+        # self.kernel.weight.requires_grad = False
+        with torch.no_grad():
+            self.kernel.weight.copy_(self.eigenvectors)
+        assert self.kernel.weight.requires_grad
+
+
+class ConcatPool(nn.Module):
+    """
+    A layer that concatenates nearby squares of an image
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        pooling_kernel_size: int,
+        pooling_kernel_stride: int,
+        pooling_kernel_padding: int,
+        d_model: int,
+        activation_kwargs: Optional[dict] = None,
+        activation: nn.Module = nn.ReLU,
+        linear_module: Type[nn.Module] = nn.Linear,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.pooling_kernel_size = pooling_kernel_size
+        self.pooling_kernel_stride = pooling_kernel_stride
+        self.pooling_kernel_padding = pooling_kernel_padding
+        self.d_model = d_model
+        self.pooling_output_size = (pooling_kernel_size**2) * self.in_channels
+        if activation_kwargs is not None:
+            self.activation = activation(**activation_kwargs)
+        else:
+            self.activation = activation()
+        self.process = nn.Sequential(
+            *[
+                nn.Unfold(
+                    pooling_kernel_size,
+                    stride=pooling_kernel_stride,
+                    padding=pooling_kernel_padding,
+                ),
+                Rearrange("N Block L -> N L Block"),
+                linear_module(
+                    self.pooling_output_size,
+                    2 * d_model if activation.__name__.endswith("GLU") else d_model,
+                ),
+                self.activation,
+            ]
+        )
+
+    def forward(self, x):
+        return self.process(x)
+
+
+class WhitenConcatProject(nn.Module):
+    """
+    A layer that whitens patches of the input image, then concatenates nearby
+        patch embeddings before rearranging the outputs into the format expected
+        by a Transformer block
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        convolution_kernel_size: int,
+        pooling_kernel_size: int,
+        pooling_kernel_stride: int,
+        pooling_kernel_padding: int,
+        d_model: int,
+        dropout=0.0,
+        activation: nn.Module = nn.ReLU,
+        activation_kwargs: Optional[dict] = None,
+        eigenvectors: torch.Tensor = None,
+        bias: bool = True,
+        linear_module: Type[nn.Module] = nn.Linear,
+    ):
+        super().__init__()
+        assert eigenvectors is not None
+        self.in_channels = in_channels
+        self.convolution_kernel_size = convolution_kernel_size
+        self.pooling_kernel_size = pooling_kernel_size
+        self.whitening_out_channels = convolution_kernel_size**2 * in_channels * 2
+        self.pooling_output_size = (
+            pooling_kernel_size**2
+        ) * self.whitening_out_channels
+        self.process = nn.Sequential(
+            *[
+                WhiteningConv(
+                    in_channels,
+                    convolution_kernel_size,
+                    eigenvectors,
+                    bias=bias,
+                    linear_module=linear_module,
+                ),
+                ConcatPool(
+                    self.whitening_out_channels,
+                    pooling_kernel_size,
+                    pooling_kernel_stride,
+                    pooling_kernel_padding,
+                    d_model,
+                    activation,
+                    activation_kwargs,
+                    linear_module,
+                ),
+                nn.Dropout(dropout),
+            ]
+        )
+
+    def forward(self, x):
+        return self.process(x)

broccoli_ml-0.1.0/broccoli/eigenpatches.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from einops import rearrange
+
+
+def eigenvectors(images: torch.Tensor, patch_size: int = 2, eps=5e-4) -> torch.Tensor:
+    """
+    Adapted from
+        https://github.com/KellerJordan/cifar10-airbench/blob/master/airbench96_faster.py
+        using https://datascienceplus.com/understanding-the-covariance-matrix/
+    """
+    with torch.no_grad():
+        unfolder = nn.Unfold(kernel_size=patch_size, stride=1)
+        patches = unfolder(images)  # (N, patch_elements, patches_per_image)
+        patches = rearrange(patches, "N elements patches -> (N patches) elements")
+        n = patches.size(0)
+        centred = patches - patches.mean(dim=1, keepdim=True)
+        covariance_matrix = (
+            centred.T @ centred
+        ) / n  # https://datascienceplus.com/understanding-the-covariance-matrix/
+        _, eigenvectors = torch.linalg.eigh(covariance_matrix)
+        return eigenvectors
+
+
+def get_eigenpatches(data: DataLoader):
+    patches = None
+    for images, _ in data:
+        eigenpatches = eigenvectors(images)
+        if patches is None:
+            patches = eigenpatches
+        else:
+            patches = 0.99 * patches + 0.01 * eigenpatches
+    return patches
+
+
+def plottable_bandpass_filters(eigenpatches: torch.Tensor, h, w, c):
+    bandpass_filters = rearrange(eigenpatches, "N (C H W) -> N H W C", C=3, H=2, W=2)
+    return bandpass_filters.detach().numpy()

broccoli_ml-0.1.0/broccoli/linear.py

@@ -0,0 +1,41 @@
+# UNDER CONSTRUCTION
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class RandomLinear(nn.Linear):
+    """ """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,  # <---- TODO: explain this
+        beta=0.1,
+        forward_looks_random=True,
+    ):
+        super().__init__(in_features, out_features, bias=False)
+        self.beta = beta
+        self.forward_looks_random = forward_looks_random
+
+    def forward(self, inputs: torch.Tensor):
+        if not self.training:
+            return F.linear(inputs, self.weight)
+        else:
+            # Initialise self.random_weights
+            random_weights = torch.empty_like(self.weight)
+            nn.init.trunc_normal_(random_weights)
+            random_weights *= self.beta
+
+            if self.forward_looks_random:
+                # Forward using a reparameterisation trick
+                a = F.linear(inputs.detach(), self.weight, self.bias)
+                b = F.linear(inputs, random_weights, bias=None)
+            else:
+                # Forward as (W_actual * input + W_random * input) + bias
+                a = F.linear(inputs, self.weight, self.bias)
+                b = F.linear(inputs, random_weights, bias=None)
+
+            return a + b

broccoli_ml-0.1.0/broccoli/tensor.py

@@ -0,0 +1,49 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class SigmaReparamTensor(nn.Module):
+    """
+    Inspired by Apple's Spectral Normed Linear Layers
+        (https://github.com/apple/ml-sigma-reparam)
+
+    """
+
+    def __init__(self, init_tensor: torch.Tensor):
+        assert init_tensor.ndim == 2
+
+        super().__init__()
+
+        self.tensor = nn.Parameter(init_tensor, requires_grad=True)
+
+        with torch.no_grad():
+            _, sigma, v_transpose = torch.linalg.svd(self.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.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 = F.normalize(
+                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)
+            self.right_singular.data.copy_(updated_right_singular)
+            rayleigh_quotient = torch.einsum(
+                "m,mn,n->",
+                approx_right_singular_transpose,
+                self.tensor,
+                updated_right_singular,
+            )
+            self.approx_spectral_norm.data.copy_(rayleigh_quotient)
+
+    def forward(self):
+        if self.training:
+            self.power_iteration()
+        return self.scale * (self.tensor / self.approx_spectral_norm)

broccoli_ml-0.1.0/broccoli/transformer.py

@@ -0,0 +1,260 @@
+import math
+from collections import OrderedDict
+from typing import Optional
+from numpy import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange, einsum, reduce, repeat
+from einops.layers.torch import Rearrange
+
+
+class MHAttention(nn.Module):
+    """
+    Multi-head self-attention using einops and custom linear layer.
+
+    Forward method assumes q, k and v have the same embedding size and k and v
+        are the same shape.
+
+    Assumes bias=False and batch_first=True, as God intended.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        n_heads,
+        dropout=0.0,
+        causal=False,
+        sequence_length=None,
+        share_kv=True,
+        linear_module: nn.Module = nn.Linear,
+    ):
+        super().__init__()
+        if causal:
+            assert sequence_length is not None
+        self.embed_dim = embed_dim
+        self.n_heads = n_heads
+        assert embed_dim % n_heads == 0
+        self.head_dim = self.embed_dim // self.n_heads
+        self.share_kv = share_kv
+        self.q_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
+        self.k_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
+        if self.share_kv:
+            self.v_proj = self.k_proj
+        else:
+            self.v_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
+        self.out_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
+        self.causal = causal
+        self.sequence_length = sequence_length
+        self.dropout = nn.Dropout(dropout)
+        if self.causal:
+            self.register_buffer(
+                "mask",
+                (
+                    torch.triu(torch.ones(sequence_length, sequence_length), diagonal=1)
+                    == 1
+                )
+                .unsqueeze(0)
+                .unsqueeze(0),
+            )
+
+    def forward(self, q, k, v):
+        query_batch_size, query_tokens, query_features = q.size()
+        key_batch_size, key_tokens, key_features = k.size()
+
+        assert k.size() == v.size()
+        assert query_features == key_features
+        assert (
+            (query_batch_size == key_batch_size)  # batch sizes are the same...
+            or query_batch_size == 1  # ... or query is broadcastable
+        )
+
+        if self.causal:
+            assert query_tokens == key_tokens
+            assert query_tokens == self.sequence_length
+
+        # Project q, k and v and divide into heads
+        q = rearrange(self.q_proj(q), "b t (h d) -> b h t d", h=self.n_heads)
+        k = rearrange(self.k_proj(k), "b t (h d) -> b h t d", h=self.n_heads)
+        if self.share_kv:
+            v = k
+        else:
+            v = rearrange(self.v_proj(v), "b t (h d) -> b h t d", h=self.n_heads)
+
+        qk_scores = q @ k.transpose(-1, -2)
+        qk_scores /= math.sqrt(self.head_dim)  # scaling
+
+        # Apply mask if causal (must come before softmax)
+        if self.causal:
+            qk_scores.masked_fill_(self.mask, float("-inf"))
+
+        qk_scores = torch.softmax(qk_scores, dim=-1)  # softmax
+        qk_scores = self.dropout(qk_scores)  # dropout must come after softmax!
+
+        output_with_heads = qk_scores @ v
+
+        output_without_heads = rearrange(output_with_heads, "b h t d -> b t (h d)")
+
+        return self.out_proj(output_without_heads)
+
+
+class TransformerBlock(nn.Module):
+    """
+    Performs LayerNorms first (as in PyTorch Transformers when norm_first=True),
+        which is also what is seen in e.g.
+        https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
+        and is recommended by https://arxiv.org/abs/2002.04745
+
+    """
+
+    def __init__(
+        self,
+        d_model,
+        n_heads,
+        mlp_ratio=4,
+        activation: nn.Module = nn.ReLU,
+        activation_kwargs: Optional[dict] = None,
+        mlp_dropout=0.0,
+        msa_dropout=0.0,
+        causal=False,
+        linear_module=nn.Linear,
+    ):
+        super().__init__()
+
+        if activation_kwargs is not None:
+            self.activation = activation(**activation_kwargs)
+        else:
+            self.activation = activation()
+
+        # Submodules for applying attention
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.attn = MHAttention(  # Handles QKV projection
+            d_model,
+            n_heads,
+            dropout=msa_dropout,
+            causal=causal,
+            linear_module=linear_module,
+        )
+
+        # Submodules for the feedforward process
+        self.ff_process = nn.Sequential(
+            OrderedDict(
+                [
+                    ("layer_norm", nn.LayerNorm(d_model)),
+                    (
+                        # up_projection is appropriate to activation
+                        "up_projection",
+                        linear_module(
+                            d_model,
+                            (
+                                2 * mlp_ratio * d_model
+                                if activation.__name__.endswith("GLU")
+                                else mlp_ratio * d_model
+                            ),
+                        ),
+                    ),
+                    # xGLU activations will halve embedding size
+                    ("activation", self.activation),
+                    ("down_projection", linear_module(mlp_ratio * d_model, d_model)),
+                    ("dropout", nn.Dropout(mlp_dropout)),
+                ]
+            )
+        )
+
+    def forward(self, x):
+        normx = self.layer_norm(x)
+        x = x + self.attn(normx, normx, normx)
+        x = x + self.ff_process(x)
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    """
+    This assumes we already get a sequence of embeddings (e.g. word or image
+        patch embeddings). It uses learned positional embeddings.
+    """
+
+    def __init__(
+        self,
+        seq_len,
+        d_model,
+        n_layers,
+        n_heads,
+        mlp_ratio=4,
+        activation: nn.Module = nn.ReLU,
+        activation_kwargs: Optional[dict] = None,
+        mlp_dropout=0.0,
+        msa_dropout=0.0,
+        stochastic_depth=0.0,
+        causal=False,
+        linear_module=nn.Linear,
+        bos_tokens=True,
+    ):
+        super().__init__()
+        self.seq_len = seq_len
+        self.n_heads = n_heads
+
+        # Initialise BOS tokens with Xavier uniform init per
+        # https://docs.pytorch.org/docs/stable/nn.init.html#torch.nn.init.xavier_uniform_
+        if bos_tokens:
+            self.bos_tokens = nn.Parameter(torch.empty(bos_tokens, d_model))
+            nn.init.normal_(self.bos_tokens, mean=0.0, std=1.0)
+            self.full_sequence_length = self.seq_len + bos_tokens
+        else:
+            self.bos_tokens = None
+            self.full_sequence_length = self.seq_len
+
+        self.d_model = d_model
+        self.positional_embedding = nn.Embedding(self.full_sequence_length, d_model)
+        self.mlp_dropout = mlp_dropout
+        self.msa_dropout = msa_dropout
+        self.stochastic_depth = stochastic_depth
+        self.blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    d_model,
+                    n_heads,
+                    mlp_ratio=mlp_ratio,
+                    activation=activation,
+                    activation_kwargs=activation_kwargs,
+                    mlp_dropout=mlp_dropout,
+                    msa_dropout=msa_dropout,
+                    causal=causal,
+                    linear_module=linear_module,
+                )
+                for _ in range(n_layers)
+            ]
+        )
+
+    def forward(self, x):
+        if self.bos_tokens is not None:
+            x = torch.cat([self.bos_tokens.expand(x.size(0), -1, -1), x], dim=1)
+        else:
+            x = x
+        x = x + self.positional_embedding(
+            torch.arange(
+                0, self.full_sequence_length, dtype=torch.long, device=x.device
+            ).unsqueeze(
+                0
+            )  # to shape (1, seq_len) to broadcast over batch
+        )
+        for block in self.blocks:
+            if (not self.training) or self.stochastic_depth == 0.0:
+                x = block(x)
+            else:  # drop out some rows from the next Transformer block operation
+                binomial = random.binomial(n=x.size(0), p=1 - self.stochastic_depth)
+                shuffle_indices = torch.randperm(x.size(0), device=x.device)
+                unshuffle_indices = torch.argsort(shuffle_indices)  # , device=x.device)
+                shuffled = x[shuffle_indices, :, :]
+                include = shuffled[:binomial, :, :]
+                exclude = shuffled[binomial:, :, :]
+                x = torch.cat([block(include), exclude])[
+                    unshuffle_indices, :, :
+                ].contiguous()
+
+        if self.bos_tokens is not None:
+            return x[:, (self.full_sequence_length - self.seq_len) :, :]
+        else:
+            return x

broccoli_ml-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[project]
+name = "broccoli-ml"
+version = "0.1.0"
+description = "Some useful Pytorch models, circa 2025"
+authors = [
+    {name = "Nicholas Bailey"}
+]
+license = {text = "MIT"}
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "torch (>=2.7.1,<3.0.0)",
+    "numpy (>=2.3.1,<3.0.0)",
+    "einops (>=0.8.1,<0.9.0)"
+]
+
+[tool.poetry]
+packages = [
+    { include = "broccoli" } # This tells Poetry that the importable package is named 'broccoli'
+]
+
+[tool.poetry.group.dev.dependencies]
+black = "^25.1.0"
+flake8 = "^7.3.0"
+pytest = "^8.4.1"
+pytest-cov = "^6.2.1"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.black]
+line-length = 88
+target-version = ['py312']
+include = '\.pyi?$'
+extend-exclude = '''
+# A regex preceded with ^/ will apply only to files and directories
+# in the root of the project.
+(
+  ^/foo.py    # exclude a file named foo.py in the root of the project
+  | .*_pb2.py  # exclude autogenerated Protocol Buffer files anywhere in the project
+)
+'''