broccoli-ml 0.26.0__py3-none-any.whl → 6.0.0__py3-none-any.whl

broccoli/cnn.py

@@ -1,300 +1,12 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn.modules.utils import _pair
-from einops import rearrange
 import math
-from typing import Type, Union, Tuple, Optional, Literal
+from typing import Union
 
 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).
-#         """
-#         _, 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
-
-
 def spatial_tuple(size: Union[int, tuple], spatial_dimensions):
     """
     Converts an integer x to `tuple([x] * spatial_dimensions)`.

broccoli/linear.py

@@ -1,5 +1,3 @@
-# UNDER CONSTRUCTION
-
 import math
 import torch
 from torch import nn
@@ -34,7 +32,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 +76,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
@@ -95,7 +95,7 @@ class AnchoredLinear(nn.Module):
         )
 
 
-class ReparamLinear(nn.Module):
+class WeightNormedLinear(nn.Module):
     """
     ...
     """
@@ -120,7 +120,8 @@ class ReparamLinear(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 +132,7 @@ class ReparamLinear(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/rope.py

@@ -27,13 +27,28 @@ SOFTWARE.
 """
 
 from __future__ import annotations
-from math import pi, log
+from math import pi
 
 import torch
-from torch.amp import autocast
+
 from torch.nn import Module
 from torch import nn, einsum, broadcast_tensors, is_tensor, tensor, Tensor
 
+# Gracefully find the best way to import autocast
+try:
+    from torch.amp import autocast as autocast_factory
+except ImportError:
+    # Fallback: For PyTorch 1.6 to 1.9
+    from torch.cuda.amp import autocast
+
+    def autocast_factory(_, enabled=True):
+        """
+        A wrapper that mimics the modern autocast signature but calls the older
+        torch.cuda.amp.autocast, ignoring the device_type argument.
+        """
+        return autocast(enabled=enabled)
+
+
 from einops import rearrange, repeat
 
 from typing import Literal
@@ -74,7 +89,7 @@ def rotate_half(x):
     return rearrange(x, "... d r -> ... (d r)")
 
 
-@autocast("cuda", enabled=False)
+@autocast_factory("cuda", enabled=False)
 def apply_rotary_emb(
     freqs, t, start_index=0, scale=1.0, seq_dim=-2, freqs_seq_dim=None
 ):
@@ -363,7 +378,7 @@ class RotaryEmbedding(Module):
         all_freqs = broadcast_tensors(*all_freqs)
         return torch.cat(all_freqs, dim=-1)
 
-    @autocast("cuda", enabled=False)
+    @autocast_factory("cuda", enabled=False)
     def forward(self, t: Tensor, seq_len: int | None = None, offset=0):
         should_cache = (
             self.cache_if_possible

broccoli/tensor.py

@@ -58,48 +58,52 @@ class SigmaReparamTensor(nn.Module):
 
 class AnchoredReparamTensor(nn.Module):
     """
-    Reparameterise a tensor as a normalised tensor of weights multiplied by a
-        learnable scaling factor.
+    Reparameterises a tensor by decoupling its magnitude and direction.
 
-    The tensor of weights is also reparameterised as the product of a learnable
-        weight tensor with the (fixed) dominant right-singular vector of the
-        weight tensor as it was initialised.
+    The direction is represented by a learnable weight tensor, normalised by the
+    Rayleigh quotient with respect to its initial dominant right-singular vector.
+    The magnitude is a separate learnable scalar.
 
-    i.e this module represents a tensor reparameterised as:
+    The reparameterization is:
 
-        W_reparam = scale * (W / ||W @ v_0||_2)
+        W_reparam = scale * (W / norm)
 
-        where v_0 is the dominant right-singular vector of the initial tensor W_init.
+    where the norm is the Rayleigh quotient uᵀWv₀, with v₀ being the dominant
+    right-singular vector of the initial tensor and u = normalize(Wv₀).
     """
 
     def __init__(self, init_tensor: torch.Tensor):
-        assert init_tensor.ndim == 2, "Input tensor must be a 2D matrix."
+        assert init_tensor.ndim == 2
+
         super().__init__()
 
-        # Use the gradboard convention of calling something nondecay_* if we should
-        # exclude it from weight decay
-        self.nondecay_weight = nn.Parameter(init_tensor.clone(), requires_grad=True)
+        self.weight = nn.Parameter(init_tensor, requires_grad=True)
 
-        # At initialization, compute the dominant right-singular vector (v_0)
-        # and store it in a non-trainable buffer.
         with torch.no_grad():
-            _, _, v_transpose = torch.linalg.svd(
-                self.nondecay_weight, full_matrices=False
-            )
-            # v_transpose[0] is the first row of V^T, which is the first right-singular vector.
-            self.register_buffer("anchor_vector", v_transpose[0])
+            _, sigma, v_transpose = torch.linalg.svd(self.weight, full_matrices=False)
 
-        initial_norm = torch.linalg.vector_norm(
-            self.nondecay_weight.mv(self.anchor_vector)
+        self.register_buffer("rayleigh_norm", sigma[:1])
+        self.register_buffer("initial_right_singular", v_transpose[0])
+        self.nondecay_scale = nn.Parameter(
+            sigma[:1].clone().detach(), requires_grad=True
         )
-        self.scale = nn.Parameter(initial_norm.clone().detach(), requires_grad=True)
 
-    def forward(self) -> torch.Tensor:
-        # Calculate the L2 norm of the matrix-vector product W @ v_0
-        norm = torch.linalg.vector_norm(self.nondecay_weight.mv(self.anchor_vector))
+    def _update_rayleigh_norm(self):
+        with torch.no_grad():
+            product = self.weight.mv(self.initial_right_singular)
+            normed_product = F.normalize(product, dim=0)
+            rayleigh_norm = torch.einsum(
+                "m,mn,n->",
+                normed_product,
+                self.weight,
+                self.initial_right_singular,
+            )
+            self.rayleigh_norm.data.copy_(rayleigh_norm)
 
-        # Return the reparameterized tensor.
-        return self.scale * (self.nondecay_weight / (norm + 1e-6))
+    def forward(self):
+        if self.training:
+            self._update_rayleigh_norm()
+        return self.nondecay_scale * (self.weight / (self.rayleigh_norm + 1e-6))
 
 
 class NormReparamTensor(nn.Module):
@@ -114,10 +118,11 @@ class NormReparamTensor(nn.Module):
 
         # Use the gradboard convention of calling something nondecay_* if we should
         # exclude it from weight decay
-        self.nondecay_weight = nn.Parameter(init_tensor.clone(), requires_grad=True)
-        self.scale = nn.Parameter(
-            torch.linalg.norm(self.nondecay_weight).clone().detach(), requires_grad=True
+        self.weight = nn.Parameter(init_tensor.clone(), requires_grad=True)
+        self.nondecay_scale = nn.Parameter(
+            torch.linalg.norm(self.weight).clone().detach(), requires_grad=True
         )
 
     def forward(self) -> torch.Tensor:
-        return self.scale * F.normalize(self.nondecay_weight)
+        norm = torch.linalg.norm(self.weight)
+        return self.nondecay_scale * (self.weight / (norm + 1e-6))