broccoli-ml 5.1.1__tar.gz

broccoli_ml-5.1.1/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-5.1.1/PKG-INFO

@@ -0,0 +1,43 @@
+Metadata-Version: 2.3
+Name: broccoli-ml
+Version: 5.1.1
+Summary: Some useful Pytorch models, circa 2025
+License: MIT
+Author: Nicholas Bailey
+Requires-Python: >=3.8
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: einops (>=0.8.1,<0.9.0)
+Description-Content-Type: text/markdown
+
+# broccoli
+
+Some useful PyTorch models, circa 2025.
+
+![broccoli](broccoli.png "Image of a rockstar made of broccoli")
+
+# Getting started
+
+You can install broccoli with
+
+```
+pip install broccoli-ml
+```
+
+PyTorch is a peer dependency of `broccoli`, which means
+  * You will need to make sure you have PyTorch installed in order to use `broccoli`
+  * PyTorch will **not** be installed automatically when you install `broccoli`
+
+We take this approach because PyTorch versioning is environment-specific and we don't know where you will want to use `broccoli`. If we automatically install PyTorch for you, there's a good chance we would get it wrong!
+
+Therefore, please also make sure you install PyTorch.
+
+# Usage examples
+
+...

broccoli_ml-5.1.1/README.md

@@ -0,0 +1,25 @@
+# broccoli
+
+Some useful PyTorch models, circa 2025.
+
+![broccoli](broccoli.png "Image of a rockstar made of broccoli")
+
+# Getting started
+
+You can install broccoli with
+
+```
+pip install broccoli-ml
+```
+
+PyTorch is a peer dependency of `broccoli`, which means
+  * You will need to make sure you have PyTorch installed in order to use `broccoli`
+  * PyTorch will **not** be installed automatically when you install `broccoli`
+
+We take this approach because PyTorch versioning is environment-specific and we don't know where you will want to use `broccoli`. If we automatically install PyTorch for you, there's a good chance we would get it wrong!
+
+Therefore, please also make sure you install PyTorch.
+
+# Usage examples
+
+...

broccoli_ml-5.1.1/broccoli/__init__.py

@@ -0,0 +1,6 @@
+from . import activation
+from . import cnn
+from . import tensor
+from . import transformer
+from . import vit
+from . import rope

broccoli_ml-5.1.1/broccoli/activation.py

@@ -0,0 +1,121 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class ReLU(nn.Module):
+    """
+    A ReLU activation function with optional clamp and leakiness.
+    """
+
+    def __init__(
+        self, clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0
+    ) -> None:
+        super().__init__()
+        self.clamp = clamp
+        self.leaky = leaky
+        self.negative_slope = negative_slope
+        self.clamp_max = clamp_max
+
+    def forward(self, x):
+        if self.leaky:
+            relu = F.leaky_relu(x, negative_slope=self.negative_slope)
+        else:
+            relu = F.relu(x)
+        if self.clamp:
+            relu = torch.clamp(relu, max=self.clamp_max)
+        return relu
+
+
+class GELU(nn.Module):
+    """
+    A GELU activation function with optional clamp.
+    """
+
+    def __init__(self, clamp=True) -> None:
+        super().__init__()
+        self.clamp = clamp
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        gelu = self.gelu(x)
+        if self.clamp:
+            gelu = torch.clamp(gelu, max=6)
+        return gelu
+
+
+class Swish(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) -> None:
+        super().__init__()
+        # Learnable parameter is called "swiglu beta" so that it is easy to find
+        #   and exclude from weight decay
+        self.swish_beta = nn.Parameter(torch.tensor([1.0]))
+
+    def forward(self, x):
+        return x * F.sigmoid(self.swish_beta * x)
+
+
+class SquaredReLU(nn.Module):
+    """
+    Squared ReLU, as shown in "ReLU^2 wins" (https://arxiv.org/abs/2402.03804) to
+      be as effective as SwiGLU for training LLMs, possibly because it can allow a
+      NN to learn multyiplication, as noted by
+      https://azizbelaweid.substack.com/p/what-is-swiglu-how-to-implement-it
+    """
+
+    def __init__(
+        self, clamp=True, leaky=True, negative_slope: float = 0.01, clamp_max=6
+    ) -> None:
+        super().__init__()
+        self.clamp = clamp
+        self.leaky = leaky
+        self.negative_slope = negative_slope
+        self.clamp_max = clamp_max
+
+    def forward(self, x):
+        if self.leaky:
+            relu = F.leaky_relu(x, negative_slope=self.negative_slope)
+        else:
+            relu = F.relu(x)
+        relu_squared = relu**2
+        if self.clamp:
+            relu_squared = torch.clamp(relu_squared, max=self.clamp_max)
+        return relu_squared
+
+
+class XGLU(nn.Module):
+    """
+    Generic Gated Linear Unit
+    """
+
+    def __init__(self, activation_module: nn.Module) -> None:
+        super().__init__()
+        self.activation = activation_module
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, value = x.chunk(2, dim=-1)
+        return self.activation(gate) * value
+
+
+def SquaredReGLU(clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0) -> XGLU:
+    """
+    Factory function that creates a GLU with a SquaredReLU activation.
+    """
+    activation_module = SquaredReLU(
+        clamp=clamp, leaky=leaky, negative_slope=negative_slope, clamp_max=clamp_max
+    )
+    return XGLU(activation_module)
+
+
+def SwiGLU() -> XGLU:
+    """
+    Factory function that creates a GLU with a Swish activation.
+    """
+    return XGLU(Swish())

broccoli_ml-5.1.1/broccoli/cnn.py

@@ -0,0 +1,157 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Union
+
+from einops.layers.torch import Rearrange
+
+
+def spatial_tuple(size: Union[int, tuple], spatial_dimensions):
+    """
+    Converts an integer x to `tuple([x] * spatial_dimensions)`.
+    Performs no operation (i.e. the identity operation) on tuples of length `spatial_dimensions`.
+    Otherwise
+    """
+    if isinstance(size, int):
+        return tuple([size] * spatial_dimensions)
+    elif isinstance(size, tuple) and (len(size) == spatial_dimensions):
+        return size
+    else:
+        raise ValueError(
+            f"For {spatial_dimensions} spatial dimensions, `size` must be "
+            f"an integer or a tuple of length {spatial_dimensions}."
+        )
+
+
+def padding_tensor(padding: tuple):
+    """
+    Converts a tuple of ints (x, y, z) into a tuple of 2-tuples,
+        like ((x, x), (y, y), (z, z)).
+
+    Performs no operation (i.e. the identity operation) on a tuple of 2-tuples.
+
+    Otherwise raises an error.
+    """
+    if all(isinstance(x, int) for x in padding):
+        return tuple([tuple([p] * 2) for p in padding])
+    elif (
+        all(isinstance(p, tuple) for p in padding)
+        and all(len(p) == 2 for p in padding)
+        and all(all(isinstance(x, int) for x in p) for p in padding)
+    ):
+        return padding
+    else:
+        raise ValueError(
+            "Padding must be a tuple of ints of a tuple of 2-tuples of ints. "
+            f"It was {padding}."
+        )
+
+
+def kd_unfold(t: torch.Tensor, kernel_size=1, stride=1, padding=0, k=2):
+    """
+    Unfold operation with k spatial dimensions.
+    Does not support dilation.
+    Only supports equal padding at top and bottom.
+    """
+    if len(t.size()[2:]) != k:
+        raise ValueError(
+            f"Input tensor size should be (N, channels, spatial dims...), so "
+            f"for k = {k}, t.size() should be a tuple of length {k + 2}."
+        )
+
+    N, C = t.size(0), t.size(1)
+
+    kernel_size = spatial_tuple(kernel_size, k)
+    stride = spatial_tuple(stride, k)
+    padding = padding_tensor(spatial_tuple(padding, k))
+
+    output = t
+    output = F.pad(output, sum(reversed(padding), ()))  # i.e. the empty tuple
+
+    for i, _ in enumerate(kernel_size):
+        output = output.unfold(i + 2, kernel_size[i], stride[i])
+
+    permutation = [0, 1] + [i + k + 2 for i in range(k)] + [i + 2 for i in range(k)]
+
+    return output.permute(*permutation).reshape(N, math.prod(kernel_size) * C, -1)
+
+
+def calculate_output_spatial_size(
+    input_spatial_size, kernel_size=1, stride=1, padding=0, dilation=0
+):
+    """
+    Calculate the output size for the spatial dimensions of a convolutional operation
+    """
+    stride = spatial_tuple(stride, len(input_spatial_size))
+
+    # Handle padding keywords that are sometimes used
+    if padding == "same":
+        output_size = ()
+        for i, in_length in enumerate(input_spatial_size):
+            output_size += (math.ceil(in_length / stride[i]),)
+        return output_size
+    elif padding == "valid":
+        padding = 0
+
+    kernel_size = spatial_tuple(kernel_size, len(input_spatial_size))
+    padding = spatial_tuple(padding, len(input_spatial_size))
+    dilation = spatial_tuple(dilation, len(input_spatial_size))
+
+    output_size = ()
+
+    for i, in_length in enumerate(input_spatial_size):
+        output_size += (
+            math.floor(
+                (in_length + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1)
+                / stride[i]
+                + 1
+            ),
+        )
+    return output_size
+
+
+class SpaceToDepth(nn.Module):
+    """
+    An operation that extracts patches from an image-like tensor and stacks
+        them channel-wise.
+    """
+
+    def __init__(self, kernel_size, stride=1, padding=0, spatial_dimensions=2):
+        """
+        Input shape should be in order (channels, spatial dims...),
+            e.g. (channels, height, width)
+        """
+
+        super().__init__()
+
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.spatial_dimensions = spatial_dimensions
+
+    def forward(self, x):
+
+        N, C, *input_spatial_size = x.size()
+
+        patches = kd_unfold(
+            x,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            k=self.spatial_dimensions,
+        )
+
+        output_spatial_size = calculate_output_spatial_size(
+            input_spatial_size=input_spatial_size,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=1,  # kd_unfold doesn't support dilation
+        )
+
+        output_channels = C * math.prod(
+            spatial_tuple(self.kernel_size, self.spatial_dimensions)
+        )
+
+        return patches.view(N, output_channels, *output_spatial_size)

broccoli_ml-5.1.1/broccoli/linear.py

@@ -0,0 +1,138 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .tensor import SigmaReparamTensor, AnchoredReparamTensor, NormReparamTensor
+
+
+class SpectralNormLinear(nn.Module):
+    """
+    Inspired by Apple's Spectral Normed Linear Layers
+        (https://github.com/apple/ml-sigma-reparam)
+    """
+
+    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
+
+        # 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:
+        weights = torch.empty(self.out_features, self.in_features)
+        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
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = SigmaReparamTensor(weights)
+
+    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 AnchoredLinear(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
+
+        # 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:
+        weights = torch.empty(self.out_features, self.in_features)
+        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
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = AnchoredReparamTensor(weights)
+
+    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"AnchoredLinear(in_features={self.in_features},"
+            f"out_features={self.out_features}, bias={self.use_bias})"
+        )
+
+
+class WeightNormedLinear(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
+
+        # 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:
+        weights = torch.empty(self.out_features, self.in_features)
+        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
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = NormReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        return (
+            f"WeightNormedLinear(in_features={self.in_features},"
+            f"out_features={self.out_features}, bias={self.use_bias})"
+        )