broccoli-ml 0.29.1__py3-none-any.whl → 10.0.1__py3-none-any.whl

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/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,6 +1,8 @@
-# UNDER CONSTRUCTION
-
 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
@@ -34,7 +36,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 +80,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
@@ -120,7 +124,8 @@ class WeightNormedLinear(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 +136,233 @@ class WeightNormedLinear(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})"
         )
+
+
+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, incoming_data=None):
+        """
+        Resets rows.
+        If incoming_data is provided, resets to the centroid (mean) of that data.
+        If not, resets to the mean of existing weights.
+        """
+        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
+
+        if incoming_data is not None:
+            target_center = self._mean_input_weights(incoming_data)
+        else:
+            target_center = self._mean_value_weights()
+
+        target_center = target_center.expand(idx_tensor.size(0), -1)
+
+        if self.xglu:
+            gate_indices = idx_tensor
+            value_indices = idx_tensor + (self.linear.out_features // 2)
+            self._update_weights(gate_indices, 0, target_center, self.optimisers)
+            self._update_weights(value_indices, 0, target_center, self.optimisers)
+        else:
+            self._update_weights(idx_tensor, 0, target_center, self.optimisers)
+
+    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_input_weights(self, input):
+        reduce_dims = list(range(input.ndim - 1))
+        data_mean = input.detach().mean(dim=reduce_dims, keepdim=True)
+
+        weights = self.linear.weight.data
+        stdv = 1.0 / math.sqrt(weights.size(1))
+        data_norm = data_mean.std() + 1e-6
+        scale_factor = stdv / data_norm
+
+        return data_mean * scale_factor
+
+    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/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