braindecode 1.8.1.dev173644701__tar.gz → 1.8.1.dev173738593__tar.gz

{braindecode-1.8.1.dev173644701/braindecode.egg-info → braindecode-1.8.1.dev173738593}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev173644701
+Version: 1.8.1.dev173738593
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -50,6 +50,7 @@ Provides-Extra: tests
 Requires-Dist: pytest<9.1; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: pytest-xdist>=3.8.0; extra == "tests"
+Requires-Dist: pytest-timeout; extra == "tests"
 Requires-Dist: codecov; extra == "tests"
 Requires-Dist: pytest_cases; extra == "tests"
 Requires-Dist: mypy; extra == "tests"

{braindecode-1.8.1.dev173644701 → braindecode-1.8.1.dev173738593}/braindecode/functional/__init__.py

@@ -1,5 +1,7 @@
+from .attention import scaled_dot_product_attention
 from .functions import (
     _get_gaussian_kernel1d,
+    apply_rotary,
     daubechies_filters,
     drop_path,
     dwt_max_level,
@@ -10,6 +12,9 @@ from .functions import (
     pairwise_iou_1d,
     plv_time,
     prefer_fft_conv,
+    rotary_angles,
+    rotary_frequencies,
+    rotate_half,
     rotate_pairs,
     safe_log,
     sinusoidal_positional_encoding,
@@ -22,9 +27,11 @@ from .initialization import (
     init_weights_,
     rescale_parameter,
 )
+from .patches import pad_or_crop_to_patches, patch_count
 
 __all__ = [
     "_get_gaussian_kernel1d",
+    "apply_rotary",
     "daubechies_filters",
     "drop_path",
     "dwt_max_level",
@@ -32,11 +39,17 @@ __all__ = [
     "hilbert_freq",
     "identity",
     "iou_1d",
+    "pad_or_crop_to_patches",
     "pairwise_iou_1d",
+    "patch_count",
     "plv_time",
     "prefer_fft_conv",
+    "rotary_angles",
+    "rotary_frequencies",
+    "rotate_half",
     "rotate_pairs",
     "safe_log",
+    "scaled_dot_product_attention",
     "sinusoidal_positional_encoding",
     "spectral_input",
     "square",

braindecode-1.8.1.dev173738593/braindecode/functional/attention.py

@@ -0,0 +1,152 @@
+"""Scaled dot-product attention with the device-specific paths in one place."""
+
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+@torch.jit.unused
+def _hpu_eager_fused_sdpa(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    scale: Optional[float] = None,
+    fused_in_lazy: bool = False,
+) -> Optional[Tensor]:
+    # Lazy mode: None, the graph compiler already fuses torch's SDPA there
+    # (unless the caller measured the fused kernel faster for its shapes).
+    from habana_frameworks.torch.utils.internal import is_lazy
+
+    if is_lazy() and not fused_in_lazy:
+        return None
+    from habana_frameworks.torch.hpex.kernels import FusedSDPA
+
+    # Strided q/k/v views (heads split off a projection) give wrong outputs and
+    # gradients at some shapes; contiguous copies give the right result.
+    query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+    return FusedSDPA.apply(query, key, value, None, 0.0, False, scale, "None", False)
+
+
+@torch.jit.unused
+def _hpu_lazy() -> bool:
+    from habana_frameworks.torch.utils.internal import is_lazy
+
+    return is_lazy()
+
+
+def _sdpa_math_dropout(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    dropout_p: float,
+    scale: Optional[float] = None,
+) -> Tensor:
+    # torch's math path for SDPA with dropout, op for op, without the pass of
+    # ``_safe_softmax`` that zeroes fully masked rows (none without a mask):
+    # same values, same dropout draws.
+    if scale is None:
+        scale = 1.0 / math.sqrt(query.size(-1))
+    half = math.sqrt(scale)
+    attn = torch.matmul(query * half, key.transpose(-2, -1) * half)
+    attn = F.dropout(attn.softmax(-1), dropout_p, True)
+    return torch.matmul(attn, value)
+
+
+def scaled_dot_product_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_mask: Optional[Tensor] = None,
+    dropout_p: float = 0.0,
+    scale: Optional[float] = None,
+    fused_in_lazy: bool = False,
+    exact: bool = False,
+) -> Tensor:
+    """``F.scaled_dot_product_attention``; Gaudi's fused kernel on HPU.
+
+    Parameters
+    ----------
+    query, key, value : torch.Tensor
+        Shape ``(..., seq, head_dim)``; ``key``/``value`` may have another
+        sequence length.
+    attn_mask : torch.Tensor, optional
+        Boolean (``True`` takes part) or additive mask, as in torch.
+    dropout_p : float, default=0.0
+        Dropout on the attention weights; pass ``0.0`` in eval mode.
+    scale : float, optional
+        Score scale; ``None`` is ``1 / sqrt(head_dim)``.
+    fused_in_lazy : bool, default=False
+        Use Gaudi's fused kernel in HPU lazy mode too.
+    exact : bool, default=False
+        Run the textbook op order instead of torch's kernels:
+        ``softmax((q @ k^T) / sqrt(d)) @ v``, or ``softmax((q @ k^T) * scale) @ v``
+        when ``scale`` is given, with dropout on the weights. Bit-identical to
+        a hand-written attention; torch's kernels round differently (its math
+        path scales ``q`` and ``k`` by ``sqrt(scale)`` each before the matmul,
+        its fused CPU kernel runs a blocked softmax), about 1e-6 apart.
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(..., seq, head_dim)``.
+
+    In HPU eager mode, torch's SDPA runs op by op (the full
+    ``(batch, heads, seq, seq)`` scores in device memory), 2-3.5x slower than
+    the bridge's fused kernel on long sequences. The fused kernel is used
+    without a mask or dropout (its dropout would draw other masks). In lazy
+    mode only with ``fused_in_lazy=True``: the graph compiler fuses torch's
+    SDPA too, faster on long sequences, slower on short ones.
+    """
+    if exact:
+        scores = torch.matmul(query, key.transpose(-2, -1))
+        if scale is None:
+            scores = scores / math.sqrt(query.size(-1))
+        else:
+            scores = scores * scale
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                scores = scores.masked_fill(~attn_mask, float("-inf"))
+            else:
+                scores = scores + attn_mask
+        weights = scores.softmax(dim=-1)
+        if dropout_p > 0.0:
+            weights = F.dropout(weights, dropout_p, True)
+        return torch.matmul(weights, value)
+    if (
+        not torch.jit.is_scripting()
+        and query.device.type == "hpu"
+        and attn_mask is None
+        and dropout_p == 0.0
+        and (query.dtype == torch.float32 or query.dtype == torch.bfloat16)
+    ):
+        out = _hpu_eager_fused_sdpa(
+            query, key, value, scale=scale, fused_in_lazy=fused_in_lazy
+        )
+        if out is not None:
+            return out
+    if (
+        not torch.jit.is_scripting()
+        and dropout_p > 0.0
+        and attn_mask is None
+        and (query.dtype == torch.float32 or query.dtype == torch.float64)
+        and (
+            query.device.type == "cpu"
+            or (query.device.type == "hpu" and not _hpu_lazy())
+        )
+    ):
+        # CPU and HPU eager run SDPA with dropout through the math path.
+        return _sdpa_math_dropout(query, key, value, dropout_p, scale)
+    if scale is None:  # torch 2.0's signature has no ``scale``
+        return F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attn_mask, dropout_p=dropout_p
+        )
+    return F.scaled_dot_product_attention(
+        query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, scale=scale
+    )

{braindecode-1.8.1.dev173644701 → braindecode-1.8.1.dev173738593}/braindecode/functional/functions.py

@@ -588,6 +588,7 @@ def sinusoidal_positional_encoding(
     dim: int,
     device: torch.device | None = None,
     dtype: torch.dtype = torch.float32,
+    base: float = 10000.0,
 ) -> torch.Tensor:
     r"""Fixed sine/cosine positional-encoding table of shape ``(n_positions, dim)``.
 
@@ -609,6 +610,9 @@ def sinusoidal_positional_encoding(
         Device of the table (default: CPU), for tables built inside ``forward``.
     dtype : torch.dtype
         Floating dtype the table is computed in.
+    base : float, default=10000.0
+        Longest wavelength, the ``10000`` of Vaswani et al.; SignalJEPA uses the
+        largest number of time samples it encodes.
 
     Returns
     -------
@@ -620,7 +624,7 @@ def sinusoidal_positional_encoding(
     position = torch.arange(n_positions, device=device).unsqueeze(1).to(dtype)
     div_term = torch.exp(
         torch.arange(0, dim_even, 2, device=device).to(dtype)
-        * (-math.log(10000.0) / dim_even)
+        * (-math.log(base) / dim_even)
     )
     pe = torch.zeros(n_positions, dim_even, device=device, dtype=dtype)
     pe[:, 0::2] = torch.sin(position * div_term)
@@ -646,3 +650,104 @@ def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
     """
     pairs = x.unflatten(-1, (-1, 2))
     return torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).flatten(-2)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate the two feature halves from ``(a, b)`` to ``(-b, a)``.
+
+    The half-split layout of GPT-NeoX / LLaMA rotary embeddings; see
+    :func:`rotate_pairs` for the interleaved layout.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Tensor of shape ``(..., features)`` with an even final dimension.
+
+    Returns
+    -------
+    torch.Tensor
+        Same shape, device and dtype as ``x``.
+    """
+    first, second = x.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def rotary_frequencies(
+    head_dim: int,
+    base: float = 10000.0,
+    device: Optional[torch.device] = None,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Inverse frequencies of a rotary embedding, ``1 / base ** (2 i / head_dim)``.
+
+    Parameters
+    ----------
+    head_dim : int
+        Even feature dimension that is rotated.
+    base : float, default=10000.0
+        Base period (``theta``).
+    device : torch.device, optional
+    dtype : torch.dtype, default=torch.float32
+        Precision of the table; the angles inherit it.
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(head_dim // 2,)``.
+    """
+    exponent = torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim
+    return 1.0 / (base**exponent)
+
+
+def rotary_angles(
+    positions: torch.Tensor, inv_freq: torch.Tensor, interleaved: bool = True
+) -> torch.Tensor:
+    """Rotation angles of each position, laid out to match the feature layout.
+
+    Parameters
+    ----------
+    positions : torch.Tensor
+        Shape ``(...,)``; same dtype as ``inv_freq`` for exact angles.
+    inv_freq : torch.Tensor
+        Shape ``(head_dim // 2,)``, from :func:`rotary_frequencies`.
+    interleaved : bool, default=True
+        ``True`` repeats each angle for its adjacent pair (:func:`rotate_pairs`);
+        ``False`` concatenates the angles twice (:func:`rotate_half`).
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(..., head_dim)``; pass ``.cos()`` and ``.sin()`` to
+        :func:`apply_rotary`.
+    """
+    angles = positions.unsqueeze(-1) * inv_freq
+    if interleaved:
+        return angles.repeat_interleave(2, dim=-1)
+    return torch.cat((angles, angles), dim=-1)
+
+
+def apply_rotary(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = True
+) -> torch.Tensor:
+    """Rotate ``x`` by the angles whose cosine and sine are given.
+
+    ``x * cos + rotate(x) * sin``, where ``rotate`` is :func:`rotate_pairs`
+    (``interleaved=True``) or :func:`rotate_half`. ``cos`` and ``sin`` must
+    broadcast to ``x``; casting is left to the caller.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Shape ``(..., head_dim)``.
+    cos, sin : torch.Tensor
+        From :func:`rotary_angles`.
+    interleaved : bool, default=True
+        Feature layout, as in :func:`rotary_angles`.
+
+    Returns
+    -------
+    torch.Tensor
+        Same shape as ``x``.
+    """
+    rotated = rotate_pairs(x) if interleaved else rotate_half(x)
+    return x * cos + rotated * sin

braindecode-1.8.1.dev173738593/braindecode/functional/patches.py

@@ -0,0 +1,105 @@
+"""Cutting a time axis into patches: how many, and padding or cropping to fit."""
+
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+import torch.nn.functional as F
+
+_ON_NON_DIVISIBLE = ("pad", "crop", "error")
+
+
+def _check_on_non_divisible(on_non_divisible: str) -> None:
+    """Raise a ``ValueError`` for an unknown ``on_non_divisible`` value."""
+    if on_non_divisible not in _ON_NON_DIVISIBLE:
+        raise ValueError(
+            "on_non_divisible must be 'pad', 'crop', or 'error', "
+            f"got {on_non_divisible!r}."
+        )
+
+
+def patch_count(
+    n_times: int, patch_size: int, on_non_divisible: str = "pad", step: int = 0
+) -> int:
+    """Number of patches cut from ``n_times`` samples.
+
+    Parameters
+    ----------
+    n_times : int
+        Length of the axis that is patched.
+    patch_size : int
+        Samples per patch.
+    on_non_divisible : {"pad", "crop", "error"}, default="pad"
+        A last, incomplete patch is zero-padded (``"pad"``), dropped
+        (``"crop"``) or raises a ``ValueError``.
+    step : int, default=0
+        Distance between patch starts; ``0`` means ``patch_size``
+        (non-overlapping).
+
+    Returns
+    -------
+    int
+        Number of patches.
+    """
+    if step <= 0:
+        step = patch_size
+    if n_times < 1:
+        raise ValueError(f"n_times must be positive, got {n_times}.")
+    if n_times < patch_size:
+        if on_non_divisible == "pad":
+            return 1
+        if on_non_divisible == "crop":
+            raise ValueError(
+                f"Need at least one full patch of {patch_size} samples, got {n_times}."
+            )
+    elif (n_times - patch_size) % step == 0:
+        return (n_times - patch_size) // step + 1
+    elif on_non_divisible != "error":
+        return (n_times - patch_size) // step + 1 + (on_non_divisible == "pad")
+    raise ValueError(
+        f"Time dimension ({n_times}) is not divisible by patch_size ({patch_size})"
+        + ("." if step == patch_size else f" with patch step {step}.")
+    )
+
+
+def pad_or_crop_to_patches(
+    x: torch.Tensor,
+    patch_size: int,
+    on_non_divisible: str = "pad",
+    step: int = 0,
+    min_patches: int = 1,
+    center: bool = False,
+) -> torch.Tensor:
+    """Right zero-pad or crop the last axis so the patches tile it.
+
+    A tiled axis is returned as is. An axis short of ``min_patches`` patches is
+    zero-padded up to that many, with half the missing ones (whole steps)
+    before it if ``center``.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Shape ``(..., n_times)``.
+    patch_size, on_non_divisible, step
+        As in :func:`patch_count`.
+    min_patches : int, default=1
+        Least number of patches on the returned axis.
+    center : bool, default=False
+        Pad half the missing patches on the left.
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(..., patch_size + (n_patches - 1) * step)``.
+    """
+    n_patches = patch_count(x.shape[-1], patch_size, on_non_divisible, step)
+    step = step if step > 0 else patch_size
+    left = 0
+    if n_patches < min_patches:
+        left = (min_patches - n_patches) // 2 * step if center else 0
+        n_patches = min_patches
+    target = patch_size + (n_patches - 1) * step
+    if target > x.shape[-1]:
+        return F.pad(x, (left, target - x.shape[-1] - left))
+    return x[..., :target]

{braindecode-1.8.1.dev173644701 → braindecode-1.8.1.dev173738593}/braindecode/models/atcnet.py

@@ -1,7 +1,6 @@
 # Authors: Cedric Rommel <cedric.rommel@inria.fr>
 #
 # License: BSD (3-clause)
-import math
 from typing import Optional
 
 import torch
@@ -10,6 +9,7 @@ from mne.utils import warn
 from torch import nn
 from torch.nn.utils.parametrize import is_parametrized, register_parametrization
 
+from braindecode.functional import scaled_dot_product_attention
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import (
     CausalConv1d,
@@ -934,20 +934,10 @@ class _MHA(nn.Module):
         K_ = torch.cat(K.split(self.head_dim, -1), 0)  # (B', S, D')
         V_ = torch.cat(V.split(self.head_dim, -1), 0)  # (B', S, D')
 
-        # Attention weights of size (num_heads * batch_size, n, m):
-        # measures how similar each pair of Q and K is.
-        W = torch.softmax(
-            Q_.bmm(K_.transpose(-2, -1)) / math.sqrt(self.head_dim),
-            -1,  # (B', D', S)
-        )  # (B', N, M)
-
-        # Multihead output (batch_size, seq_len, dim):
-        # weighted sum of V where a value gets more weight if its corresponding
-        # key has larger dot product with the query.
+        # Multihead output (batch_size, seq_len, dim): softmax(Q K^T / sqrt(d)) V
+        # per head, heads merged back along the feature axis.
         H = torch.cat(
-            (W.bmm(V_)).split(  # (B', S, S)  # (B', S, D')
-                batch_size, 0
-            ),  # [(B, S, D')] * num_heads
+            scaled_dot_product_attention(Q_, K_, V_, exact=True).split(batch_size, 0),
             -1,
         )  # (B, S, D)
 

{braindecode-1.8.1.dev173644701 → braindecode-1.8.1.dev173738593}/braindecode/models/attn_sleep.py

@@ -27,15 +27,14 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import math
 import warnings
 from copy import deepcopy
 from numbers import Integral
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 
+from braindecode.functional import scaled_dot_product_attention
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import CausalConv1d
 
@@ -411,14 +410,6 @@ class _MRCNN(nn.Module):
 ##########################################################################################
 
 
-def _attention(query: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
-    """Scaled dot product attention weights."""
-    # d_k - dimension of the query and key vectors
-    d_k = query.size(-1)
-    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
-    return F.softmax(scores, dim=-1)  # attention weights
-
-
 class _MultiHeadedAttention(nn.Module):
     def __init__(self, h, d_model, after_reduced_cnn_size, dropout=0.1):
         """Take in model size and number of heads."""
@@ -464,9 +455,14 @@ class _MultiHeadedAttention(nn.Module):
             .transpose(1, 2)
         )
 
-        # apply dropout to the *weights*, then the weighted sum
-        attn = self.dropout(_attention(query, key))
-        x = torch.matmul(attn, value)
+        # exact: the reference op order, same outputs as its hand-written attention
+        x = scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            exact=True,
+        )
 
         # merge heads and project
         x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_per_head)

{braindecode-1.8.1.dev173644701 → braindecode-1.8.1.dev173738593}/braindecode/models/axon.py

@@ -11,10 +11,12 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
-from braindecode.functional import sinusoidal_positional_encoding
+from braindecode.functional import (
+    scaled_dot_product_attention,
+    sinusoidal_positional_encoding,
+)
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import warn_if_sfreq_differs
-from braindecode.modules.attention import _scaled_dot_product_attention
 from braindecode.modules.channels import _resolve
 
 
@@ -468,7 +470,7 @@ class _AxisAttention(nn.Module):
             .unbind(0)
         )
         drop = self.att_drop_prob if self.training else 0.0
-        attended = _scaled_dot_product_attention(
+        attended = scaled_dot_product_attention(
             query, key, value, attn_mask=mask, dropout_p=drop
         )
         return attended.transpose(1, 2).reshape(n_seq, length, -1)

{braindecode-1.8.1.dev173644701 → braindecode-1.8.1.dev173738593}/braindecode/models/barista.py

@@ -11,12 +11,18 @@ import torch
 from einops.layers.torch import Rearrange
 from torch import nn
 
-from braindecode.functional import init_weights_
+from braindecode.functional import (
+    apply_rotary,
+    init_weights_,
+    patch_count,
+    rotary_angles,
+    rotary_frequencies,
+    scaled_dot_product_attention,
+)
+from braindecode.functional.patches import _check_on_non_divisible
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import extract_channel_locations_from_chs_info
 from braindecode.modules import GatedLinearUnit, PatchTokenizer
-from braindecode.modules.attention import _scaled_dot_product_attention
-from braindecode.modules.blocks import _check_on_non_divisible, _patch_count
 
 # Weight initialization of the reference implementation.
 _init_weights = partial(
@@ -287,7 +293,7 @@ class BaRISTA(EEGModuleMixin, nn.Module, license="other"):
         # ceil(n_times / patch_size) with "pad"; floor with "crop" (the
         # reference: trailing samples dropped, raises below one patch); exact
         # with "error" (raises otherwise). The learned pooling is sized to it.
-        self.n_patches = _patch_count(self.n_times, patch_size, on_non_divisible)
+        self.n_patches = patch_count(self.n_times, patch_size, on_non_divisible)
 
         self.patch_tokenizer = PatchTokenizer(
             patch_size=patch_size,
@@ -609,7 +615,7 @@ class _Transformer(nn.Module):
     ):
         super().__init__()
         head_dim = d_model // num_heads
-        inv_freq = 1.0 / (10000.0 ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        inv_freq = rotary_frequencies(head_dim)
         self.inv_freq: torch.Tensor
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.layers = nn.ModuleList(
@@ -633,8 +639,8 @@ class _Transformer(nn.Module):
         positions = torch.arange(
             n_patches, device=self.inv_freq.device, dtype=self.inv_freq.dtype
         )
-        angles = torch.outer(positions, self.inv_freq)
-        angles = torch.cat((angles, angles), dim=-1).repeat_interleave(n_chans, dim=0)
+        angles = rotary_angles(positions, self.inv_freq, interleaved=False)
+        angles = angles.repeat_interleave(n_chans, dim=0)
         cos = angles.cos()[None, None]
         sin = angles.sin()[None, None]
         for layer in self.layers:
@@ -698,20 +704,13 @@ class _RotarySelfAttention(nn.Module):
         key = self.split_heads(key)
         value = self.split_heads(value)
 
-        query = _apply_rotary(query, cos, sin)
-        key = _apply_rotary(key, cos, sin)
+        cos, sin = cos.to(query.dtype), sin.to(query.dtype)
+        query = apply_rotary(query, cos, sin, interleaved=False)
+        key = apply_rotary(key, cos, sin, interleaved=False)
 
         # Dropout falls on the attention output rather than on the weights,
         # through dropout_p, because the reference disables the latter outright.
-        attention = _scaled_dot_product_attention(query, key, value)
+        attention = scaled_dot_product_attention(query, key, value)
         attention = self.dropout(attention)
         attention = self.merge_heads(attention)
         return self.o_proj(attention)
-
-
-def _apply_rotary(
-    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
-) -> torch.Tensor:
-    """Apply the reference's half-split (not interleaved) rotary embedding."""
-    first, second = x.chunk(2, dim=-1)
-    return x * cos.to(x.dtype) + torch.cat((-second, first), dim=-1) * sin.to(x.dtype)