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

broccoli/activation.py

@@ -1,10 +1,50 @@
 import torch
 from torch import nn
 from torch.nn import functional as F
-from einops import rearrange
 
 
-class SwiGLU(nn.Module):
+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.
@@ -12,16 +52,14 @@ class SwiGLU(nn.Module):
     Halves the incoming parameter count, which should be scaled up before input.
     """
 
-    def __init__(self, linear_module: nn.Module = nn.Linear) -> None:
+    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.swiglu_beta = nn.Parameter(torch.tensor([0.0]))
+        self.swish_beta = nn.Parameter(torch.tensor([1.0]))
 
     def forward(self, x):
-        gate, value = rearrange(x, "... (split c) -> split ... c", split=2)
-        beta_swish = gate * F.sigmoid(self.swiglu_beta * gate)
-        return beta_swish * value
+        return x * F.sigmoid(self.swish_beta * x)
 
 
 class SquaredReLU(nn.Module):
@@ -32,54 +70,52 @@ class SquaredReLU(nn.Module):
       https://azizbelaweid.substack.com/p/what-is-swiglu-how-to-implement-it
     """
 
-    def __init__(self, clamp=True, leaky=True) -> None:
+    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):
-        relu_squared = F.relu(x) ** 2
-        if self.clamp:
-            relu_squared = torch.clamp(relu_squared, max=6)
         if self.leaky:
-            relu_squared = relu_squared + 0.1 * x
+            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 ReLU(nn.Module):
+class XGLU(nn.Module):
     """
-    A ReLU activation function with optional clamp and leakiness.
+    Generic Gated Linear Unit
     """
 
-    def __init__(self, clamp=True, leaky=True) -> None:
+    def __init__(self, activation_module: nn.Module) -> None:
         super().__init__()
-        self.clamp = clamp
-        self.leaky = leaky
+        self.activation = activation_module
 
-    def forward(self, x):
-        relu = F.relu(x)
-        if self.clamp:
-            relu = torch.clamp(relu, max=6)
-        if self.leaky:
-            relu = relu + 0.1 * x
-        return relu
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, value = x.chunk(2, dim=-1)
+        return self.activation(gate) * value
 
 
-class GELU(nn.Module):
+def SquaredReGLU(clamp=True, leaky=True, negative_slope=0.01, clamp_max=6.0) -> XGLU:
     """
-    A ReLU activation function with optional clamp and leakiness.
+    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 __init__(self, clamp=True, leaky=True) -> None:
-        super().__init__()
-        self.clamp = clamp
-        self.leaky = leaky
-        self.gelu = nn.GELU()
 
-    def forward(self, x):
-        gelu = self.gelu(x)
-        if self.clamp:
-            gelu = torch.clamp(gelu, max=6)
-        if self.leaky:
-            gelu = gelu + 0.1 * x
-        return gelu
+def SwiGLU() -> XGLU:
+    """
+    Factory function that creates a GLU with a Swish activation.
+    """
+    return XGLU(Swish())

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,41 +1,138 @@
-# UNDER CONSTRUCTION
-
+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
 
-class RandomLinear(nn.Linear):
-    """ """
-
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        bias: bool = False,  # <---- TODO: explain this
-        beta=0.1,
-        forward_looks_random=True,
-    ):
-        super().__init__(in_features, out_features, bias=False)
-        self.beta = beta
-        self.forward_looks_random = forward_looks_random
-
-    def forward(self, inputs: torch.Tensor):
-        if not self.training:
-            return F.linear(inputs, self.weight)
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
         else:
-            # Initialise self.random_weights
-            random_weights = torch.empty_like(self.weight)
-            nn.init.trunc_normal_(random_weights)
-            random_weights *= self.beta
-
-            if self.forward_looks_random:
-                # Forward using a reparameterisation trick
-                a = F.linear(inputs.detach(), self.weight, self.bias)
-                b = F.linear(inputs, random_weights, bias=None)
-            else:
-                # Forward as (W_actual * input + W_random * input) + bias
-                a = F.linear(inputs, self.weight, self.bias)
-                b = F.linear(inputs, random_weights, bias=None)
-
-            return a + b
+            # 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})"
+        )

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