broccoli-ml 7.0.0__tar.gz → 9.5.1__tar.gz

{broccoli_ml-7.0.0 → broccoli_ml-9.5.1}/PKG-INFO

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

{broccoli_ml-7.0.0 → broccoli_ml-9.5.1}/broccoli/activation.py

@@ -46,10 +46,7 @@ class GELU(nn.Module):
 
 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.
+    Implementation of (beta) Swish
     """
 
     def __init__(self) -> None:

broccoli_ml-9.5.1/broccoli/linear.py

@@ -0,0 +1,352 @@
+import math
+import random
+import warnings
+from typing import Union, List, Iterable
+
+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})"
+        )
+
+
+class RecyclingLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        row_recycling_rate: float = 0.0,
+        column_recycling_rate: float = 0.0,
+        adaptive=False,
+        xglu=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.bias = bias
+        self.xglu = xglu
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        self.row_recycling_rate = row_recycling_rate
+        self.column_recycling_rate = column_recycling_rate
+        self.adaptive = adaptive
+        self.optimisers = []
+        self.initial_learning_rates = []
+        self._warned_about_registration = False
+
+    def register_optimiser(self, optimiser: torch.optim.Optimizer):
+        self.optimisers.append(optimiser)
+        self.initial_learning_rates.append(self._get_learning_rate(optimiser))
+        if self.initial_learning_rates[-1] == 0.0:
+            warnings.warn(
+                "Learning rate of registered optimiser was 0.0 - make sure "
+                "you haven't initialised a scheduler before registering the "
+                "optimiser",
+                stacklevel=2,
+            )
+
+    def _get_learning_rate(self, optimiser: torch.optim.Optimizer):
+        for group in optimiser.param_groups:
+            for param in group["params"]:
+                if param is self.linear.weight:
+                    return group["lr"]
+
+    def _get_multiplier(self):
+        if not self.adaptive or not self.optimisers:
+            return 1.0
+        else:
+            init = self.initial_learning_rates
+            current = [self._get_learning_rate(o) for o in self.optimisers]
+            pairs = zip(current, init, strict=True)
+            multipliers = [a / b for a, b in pairs if b != 0.0]
+            return min(multipliers) if multipliers else 0.0
+
+    def reset_rows(self, indices):
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.size(0):
+            value_indices = indices
+            centred_value_weights = self._mean_value_weights()
+            centred_value_weights = centred_value_weights.expand(indices.size(0), -1)
+            if self.xglu:
+                gate_indices = indices
+                value_indices = indices + (self.linear.out_features // 2)
+                centred_gate_weights = self._mean_gate_weights()
+                centred_gate_weights = centred_gate_weights.expand(indices.size(0), -1)
+                self._update_weights(
+                    gate_indices, 0, centred_gate_weights, self.optimisers  # dim
+                )
+            self._update_weights(
+                value_indices, 0, centred_value_weights, self.optimisers
+            )
+        else:
+            return
+
+    def reset_columns(self, indices):
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.size(0):
+            random_weights = self._random_weights(
+                self.linear.weight.size(0), indices.size(0)
+            )
+            # Make random col weights quiet so they don't introduce loud noise...
+            # ...but not so quiet that FP16 zeros them and ruins symmetry breaking!
+            random_weights *= 0.1
+            self._update_weights(indices, 1, random_weights, self.optimisers)  # dim
+        else:
+            return
+
+    def forward(self, x):
+        if self.training and self.optimisers:
+            self.reset_rows(self.get_reset_indices(0))
+            self.reset_columns(self.get_reset_indices(1))
+        elif self.training and not self._warned_about_registration:
+            warnings.warn(
+                "RecyclingLinear: No optimiser registered. Recycling disabled.",
+                stacklevel=2,
+            )
+            self._warned_about_registration = True
+
+        return self.linear(x)
+
+    def get_reset_indices(self, dim):
+        base_rate = self.row_recycling_rate if dim == 0 else self.column_recycling_rate
+        p = base_rate * self._get_multiplier()
+        if dim == 0:
+            if self.xglu:
+                sample_space = self.linear.out_features // 2
+            else:
+                sample_space = self.linear.out_features
+        elif dim == 1:
+            sample_space = self.linear.in_features
+        else:
+            raise ValueError("`dim` must be 0 or 1")
+
+        # Sample the indices
+        probs = torch.rand(sample_space, device=self.linear.weight.device)
+        mask = probs < p
+        if mask.any():
+            return torch.nonzero(mask).squeeze(-1)
+        else:
+            return torch.tensor([], dtype=torch.long, device=self.linear.weight.device)
+
+    def _random_weights(self, rows, columns):
+        device = self.linear.weight.device
+        weights = self.linear.weight.data
+        stdv = 1.0 / math.sqrt(weights.size(1))
+        random_weights = torch.rand(rows, columns, device=device)
+        random_weights -= 0.5  # Range [-0.5, +0.5]
+        random_weights *= 2.0 * stdv  # Range [-stdv, +stdv]
+        return random_weights
+
+    def _mean_value_weights(self):
+        """
+        Only used when self.xglu
+        """
+        weights = self.linear.weight.data
+        rows = weights.size(0)
+        if self.xglu:
+            return self.linear.weight[int(rows / 2) :].data.mean(dim=0, keepdim=True)
+        else:
+            return self.linear.weight.data.mean(dim=0, keepdim=True)
+
+    def _mean_gate_weights(self):
+        """
+        Only used when self.xglu
+        """
+        weights = self.linear.weight.data
+        rows = weights.size(0)
+        return self.linear.weight[: int(rows / 2)].data.mean(dim=0, keepdim=True)
+
+    def _update_weights(
+        self,
+        indices: Iterable[int],
+        dim: int,
+        data: torch.Tensor,
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            if dim == 0:
+                self.linear.weight.data[idx_tensor] = data
+            elif dim == 1:
+                self.linear.weight.data[:, idx_tensor] = data
+            else:
+                raise ValueError("`dim` must be 0 or 1")
+
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=dim)
+
+    def _reset_optim_state(self, param, idx_tensor, optimisers, dim):
+        """
+        Zeroes out the optimizer state for the given indices in a single operation.
+        """
+        for optimiser in optimisers:
+            if param not in optimiser.state:
+                continue
+            state = optimiser.state[param]
+
+            for _, buffer in state.items():
+                if torch.is_tensor(buffer) and buffer.shape == param.shape:
+                    # Vectorized zeroing
+                    if dim == 0:
+                        buffer[idx_tensor] = 0.0
+                    else:
+                        buffer[:, idx_tensor] = 0.0

{broccoli_ml-7.0.0 → broccoli_ml-9.5.1}/broccoli/transformer.py

@@ -325,6 +325,8 @@ class FeedforwardBlock(nn.Module):
         activation=nn.ReLU,
         activation_kwargs=None,
         dropout=0.0,
+        inner_dropout=None,
+        outer_dropout=None,
         linear_module_up=nn.Linear,
         linear_module_down=nn.Linear,
         pre_norm=True,
@@ -338,6 +340,7 @@ class FeedforwardBlock(nn.Module):
         self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
+        self.xglu = activation.__name__.endswith("GLU")
 
         if self.residual_path and (output_features < input_features):
             raise ValueError(
@@ -354,29 +357,63 @@ class FeedforwardBlock(nn.Module):
         else:
             self.activation = activation()
 
-        self.dropout = nn.Dropout(dropout)
+        self.inner_dropout = nn.Dropout(
+            inner_dropout if inner_dropout is not None else dropout
+        )
+        self.outer_dropout = nn.Dropout(
+            outer_dropout if outer_dropout is not None else dropout
+        )
 
         self.max_features = (
-            2 * ratio * output_features
-            if activation.__name__.endswith("GLU")
-            else ratio * output_features
+            2 * ratio * output_features if self.xglu else ratio * output_features
         )
 
+        self.linear_in = linear_module_up(input_features, self.max_features)
+        self.linear_out = linear_module_down(ratio * output_features, output_features)
+
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
-                linear_module_up(input_features, self.max_features),
+                self.linear_in,
                 self.activation,
+                self.inner_dropout,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                linear_module_down(ratio * output_features, output_features),
-                self.dropout,
+                self.linear_out,
+                self.outer_dropout,
             ]
         )
 
+        self.recycling_enabled = False
+        if hasattr(self.linear_in, "row_recycling_rate") and hasattr(
+            self.linear_out, "column_recycling_rate"
+        ):
+            self.recycling_enabled = True
+            self.master_recycling_rate = self.linear_in.row_recycling_rate
+            self.linear_in.row_recycling_rate = 0.0
+            self.linear_out.column_recycling_rate = 0.0
+            if (
+                hasattr(self.linear_in, "column_recycling_rate")
+                and self.linear_in.column_recycling_rate > 0
+            ) or (
+                hasattr(self.linear_out, "row_recycling_rate")
+                and self.linear_out.row_recycling_rate > 0
+            ):
+                raise NotImplementedError(
+                    "At the moment this layer can only support recycling linear "
+                    "layers if the in layer resets only rows and the out layer "
+                    "resets only columns."
+                )
+
         self.reset_parameters()
 
     def forward(self, x):
 
+        # Recycle weights if using recycling linear layers
+        if self.training and self.recycling_enabled:
+            indices = self.linear_out.get_reset_indices(1)
+            self.linear_in.reset_rows(indices)
+            self.linear_out.reset_columns(indices)
+
         if self.checkpoint:
             processed = checkpoint(self.process, x, use_reentrant=False)
         else:
@@ -422,7 +459,9 @@ class TransformerBlock(nn.Module):
         ff_linear_module_up=None,
         ff_linear_module_down=None,
         msa_scaling="d",
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
         causal=False,
@@ -484,7 +523,9 @@ class TransformerBlock(nn.Module):
             d_model,
             activation=activation,
             activation_kwargs=activation_kwargs,
-            dropout=mlp_dropout,
+            dropout=ff_dropout,
+            inner_dropout=ff_inner_dropout,
+            outer_dropout=ff_outer_dropout,
             linear_module_up=(
                 ff_linear_module_up
                 if ff_linear_module_up is not None
@@ -567,7 +608,9 @@ class TransformerEncoder(nn.Module):
         activation_kwargs: Optional[dict] = None,
         ff_linear_module_up=None,
         ff_linear_module_down=None,
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         stochastic_depth=0.0,
         causal=False,
@@ -629,7 +672,7 @@ class TransformerEncoder(nn.Module):
         else:
             self.absolute_position_embedding = None
 
-        self.mlp_dropout = mlp_dropout
+        self.mlp_dropout = ff_dropout
         self.msa_dropout = msa_dropout
         self.stochastic_depth = stochastic_depth
 
@@ -658,7 +701,9 @@ class TransformerEncoder(nn.Module):
                     ff_linear_module_up=ff_linear_module_up,
                     ff_linear_module_down=ff_linear_module_down,
                     msa_scaling=msa_scaling,
-                    mlp_dropout=mlp_dropout,
+                    ff_dropout=ff_dropout,
+                    ff_inner_dropout=ff_inner_dropout,
+                    ff_outer_dropout=ff_outer_dropout,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],
                     causal=causal,

{broccoli_ml-7.0.0 → broccoli_ml-9.5.1}/broccoli/vit.py

@@ -161,7 +161,9 @@ class ViTEncoder(nn.Module):
         transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
-        transformer_initial_ff_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -178,7 +180,9 @@ class ViTEncoder(nn.Module):
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -333,7 +337,9 @@ class ViTEncoder(nn.Module):
                 ff_linear_module_up=transformer_ff_linear_module_up,
                 ff_linear_module_down=transformer_ff_linear_module_down,
                 msa_scaling=transformer_msa_scaling,
-                mlp_dropout=transformer_mlp_dropout,
+                ff_dropout=transformer_ff_dropout,
+                ff_inner_dropout=transformer_ff_inner_dropout,
+                ff_outer_dropout=transformer_ff_outer_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
                 causal=False,
@@ -357,9 +363,21 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=(
                     # First truthy assigned value
-                    transformer_initial_ff_mlp_dropout
-                    if transformer_initial_ff_mlp_dropout is not None
-                    else transformer_mlp_dropout
+                    transformer_initial_ff_dropout
+                    if transformer_initial_ff_dropout is not None
+                    else transformer_ff_dropout
+                ),
+                inner_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_inner_dropout
+                    if transformer_initial_ff_inner_dropout is not None
+                    else transformer_ff_inner_dropout
+                ),
+                outer_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_outer_dropout
+                    if transformer_initial_ff_outer_dropout is not None
+                    else transformer_ff_outer_dropout
                 ),
                 linear_module_up=(
                     # First truthy assigned value
@@ -441,7 +459,9 @@ class ViT(nn.Module):
         transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
-        transformer_initial_ff_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -458,7 +478,9 @@ class ViT(nn.Module):
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -508,7 +530,9 @@ class ViT(nn.Module):
             transformer_initial_ff_residual_path=transformer_initial_ff_residual_path,
             transformer_initial_ff_linear_module_up=transformer_initial_ff_linear_module_up,
             transformer_initial_ff_linear_module_down=transformer_initial_ff_linear_module_down,
-            transformer_initial_ff_mlp_dropout=transformer_initial_ff_mlp_dropout,
+            transformer_initial_ff_dropout=transformer_initial_ff_dropout,
+            transformer_initial_ff_inner_dropout=transformer_initial_ff_inner_dropout,
+            transformer_initial_ff_outer_dropout=transformer_initial_ff_outer_dropout,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
             transformer_post_norm=transformer_post_norm,
@@ -525,7 +549,9 @@ class ViT(nn.Module):
             transformer_ff_linear_module_up=transformer_ff_linear_module_up,
             transformer_ff_linear_module_down=transformer_ff_linear_module_down,
             transformer_msa_scaling=transformer_msa_scaling,
-            transformer_mlp_dropout=transformer_mlp_dropout,
+            transformer_ff_dropout=transformer_ff_dropout,
+            transformer_ff_inner_dropout=transformer_ff_inner_dropout,
+            transformer_ff_outer_dropout=transformer_ff_outer_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,

{broccoli_ml-7.0.0 → broccoli_ml-9.5.1}/pyproject.toml

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

broccoli_ml-7.0.0/broccoli/linear.py

@@ -1,138 +0,0 @@
-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})"
-        )