Braindecode 1.8.1.dev170227516__tar.gz → 1.8.1.dev170248173__tar.gz

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/NOTICE.txt

@@ -42,7 +42,6 @@ The following components are licensed under the MIT License:
 - `braindecode/models/medformer.py`
 - `braindecode/models/tcformer.py`
 - `braindecode/models/ifnet.py`
-- `braindecode/models/mirepnet.py`
 
 ### Apache-2.0 Licensed Files
 
@@ -51,16 +50,6 @@ The following components are licensed under the Apache License 2.0:
 - `braindecode/models/mvpformer.py`
 - `braindecode/models/zuna.py`
 - `braindecode/models/luna.py`
-- `braindecode/models/brant.py`
-- `braindecode/models/diver1.py`
-- `braindecode/models/mscformer.py`
-
-### USC Academic License Files
-
-[License](https://github.com/ShanechiLab/BaRISTA/blob/master/LICENSE.md)
-
-- `braindecode/models/barista.py`
-- `docs/_static/model/barista.jpeg`
 
 ## License Links
 
@@ -69,7 +58,6 @@ The following components are licensed under the Apache License 2.0:
 - [CC BY-NC-SA 4.0 License](https://creativecommons.org/licenses/by-nc-sa/4.0/)
 - [MIT License](https://opensource.org/licenses/MIT)
 - [Apache-2.0 License](https://www.apache.org/licenses/LICENSE-2.0)
-- [USC Academic License](https://github.com/ShanechiLab/BaRISTA/blob/master/LICENSE.md)
 
 ## Note
 

{braindecode-1.8.1.dev170227516/braindecode.egg-info → braindecode-1.8.1.dev170248173}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev170227516
+Version: 1.8.1.dev170248173
 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>
@@ -22,8 +22,8 @@ Requires-Python: >=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 License-File: NOTICE.txt
-Requires-Dist: torch>=2.4
-Requires-Dist: torchaudio>=2.4
+Requires-Dist: torch>=2.0
+Requires-Dist: torchaudio>=2.0
 Requires-Dist: mne>=1.11.0
 Requires-Dist: mne_bids>=0.18
 Requires-Dist: h5py

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/datasets/base.py

@@ -5,7 +5,6 @@
 #          Simon Brandt <simonbrandt@protonmail.com>
 #          David Sabbagh <dav.sabbagh@gmail.com>
 #          Robin Schirrmeister <robintibor@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -1399,9 +1398,6 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
                 df = ds._windows.metadata
             else:
                 df = ds.metadata
-            if isinstance(df, pd.DataFrame):
-                # copy so the description columns stay out of the dataset
-                df = df.copy()
             for k, v in ds.description.items():
                 df[k] = v
             all_dfs.append(df)

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/datautil/serialization.py

@@ -1,7 +1,6 @@
 """Convenience functions for storing and loading of windows datasets."""
 
 # Authors: Lukas Gemein <l.gemein@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -271,9 +270,7 @@ def _load_parallel(path, i, preload, is_raw, has_stored_windows):
     else:
         window_kwargs = _load_kwargs_json("window_kwargs", sub_dir)
         windows_ds_kwargs = [
-            kwargs[1]
-            for kwargs in window_kwargs
-            if kwargs[0] in ("WindowsDataset", "EEGWindowsDataset")
+            kwargs[1] for kwargs in window_kwargs if kwargs[0] == "WindowsDataset"
         ]
         windows_ds_kwargs = windows_ds_kwargs[0] if len(windows_ds_kwargs) == 1 else {}
         if is_raw:

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/functional/__init__.py

@@ -8,7 +8,6 @@ from .functions import (
     iou_1d,
     pairwise_iou_1d,
     plv_time,
-    rotate_pairs,
     safe_log,
     sinusoidal_positional_encoding,
     square,
@@ -26,7 +25,6 @@ __all__ = [
     "iou_1d",
     "pairwise_iou_1d",
     "plv_time",
-    "rotate_pairs",
     "safe_log",
     "sinusoidal_positional_encoding",
     "square",

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/functional/functions.py

@@ -12,8 +12,6 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 
-from . import _real_dft
-
 
 def square(x):
     return x * x
@@ -167,40 +165,18 @@ def hilbert_freq(x: torch.Tensor, forward_fourier: bool = True) -> torch.Tensor:
     https://github.com/scipy/scipy/blob/v1.14.1/scipy/signal/_signaltools.py#L2287-L2394
 
     """
-    if _real_dft.needs_real_dft(x):
-        return _real_dft.hilbert_freq_real(x, forward_fourier).to(x.dtype)
-
-    input_dtype = x.dtype
-    # ``view_as_complex`` does not accept bfloat16 real/imaginary pairs.  The
-    # HPU path can produce bfloat16 Fourier coefficients under autocast, so do
-    # the complex-valued part in float32 and restore the real-valued contract
-    # at the boundary.  Promoting before the FFT also covers backends that do
-    # not implement bfloat16 FFT kernels.
-    if input_dtype == torch.bfloat16:
-        x = x.float()
-
     if forward_fourier:
-        seq_len = x.shape[-1]
         x = torch.fft.rfft(x, norm=None, dim=-1)
         x = torch.view_as_real(x)
-    else:
-        # The one-sided spectrum does not record whether the signal length
-        # was odd or even, so assume an even length here.
-        seq_len = 2 * (x.shape[-2] - 1)
     x = x * 2.0
     x[..., 0, :] = x[..., 0, :] / 2.0  # Don't multiply the DC-term by 2
-    if forward_fourier and seq_len % 2 == 0:
-        # Nor the Nyquist term, which only exists for even lengths
-        x[..., -1, :] = x[..., -1, :] / 2.0
     x = F.pad(
-        x, [0, 0, 0, seq_len - x.shape[-2]]
+        x, [0, 0, 0, x.shape[-2] - 2]
     )  # Fill Fourier coefficients to retain shape
     x = torch.view_as_complex(x)
     x = torch.fft.ifft(x, norm=None, dim=-1)  # returns complex signal
     x = torch.view_as_real(x)
 
-    if input_dtype == torch.bfloat16:
-        x = x.to(input_dtype)
     return x
 
 
@@ -470,21 +446,3 @@ def sinusoidal_positional_encoding(n_positions: int, dim: int) -> torch.Tensor:
     # ``.contiguous()`` so an odd-``dim`` truncation owns tight storage -- a
     # non-contiguous view over the padded width breaks safetensors buffer saving.
     return pe[:, :dim].contiguous()
-
-
-def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
-    """Rotate adjacent feature pairs from ``(a, b)`` to ``(-b, a)``.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        Tensor of shape ``(..., features)`` with an even final dimension.
-
-    Returns
-    -------
-    torch.Tensor
-        Same shape, device and dtype as ``x``. Frequencies, positional indexing
-        and mixed-precision policy are left to the caller.
-    """
-    pairs = x.unflatten(-1, (-1, 2))
-    return torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).flatten(-2)

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/__init__.py

@@ -3,13 +3,10 @@
 from .atcnet import ATCNet
 from .attentionbasenet import AttentionBaseNet
 from .attn_sleep import AttnSleep
-from .barista import BaRISTA
 from .base import EEGModuleMixin
 from .bendr import BENDR, InterpolatedBENDR
 from .biot import BIOT, InterpolatedBIOT
-from .brainbert import BrainBERT
 from .brainmodule import BrainModule
-from .brant import Brant
 from .cbramod import CBraMod
 from .codebrain import CodeBrain
 from .contrawr import ContraWR
@@ -18,7 +15,6 @@ from .dance import DANCE
 from .deep4 import Deep4Net
 from .deepsleepnet import DeepSleepNet
 from .dgcnn import DGCNN
-from .diver1 import DIVER1
 from .eegconformer import EEGConformer
 from .eegdino import EEGDINO
 from .eeginception_erp import EEGInceptionERP
@@ -42,8 +38,6 @@ from .labram import InterpolatedLaBraM, Labram
 from .luna import LUNA
 from .medformer import MEDFormer
 from .meta_neuromotor import MetaNeuromotorHand
-from .mirepnet import MIRepNet
-from .mscformer import MSCFormer
 from .msvtnet import MSVTNet
 from .mvpformer import MVPFormer
 from .neuropose import NeuroPose
@@ -90,11 +84,9 @@ __all__ = [
     "ATCNet",
     "AttnSleep",
     "AttentionBaseNet",
-    "BaRISTA",
     "EEGModuleMixin",
     "BIOT",
     "BENDR",
-    "BrainBERT",
     "CBraMod",
     "CodeBrain",
     "ContraWR",
@@ -102,9 +94,7 @@ __all__ = [
     "DANCE",
     "Deep4Net",
     "DeepSleepNet",
-    "DIVER1",
     "BrainModule",
-    "Brant",
     "EEGConformer",
     "EEGDINO",
     "EEGPT",
@@ -139,8 +129,6 @@ __all__ = [
     "extract_channel_locations_from_chs_info",
     "positions_from_chs_info",
     "MEDFormer",
-    "MIRepNet",
-    "MSCFormer",
     "MSVTNet",
     "MVPFormer",
     "PBT",

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/base.py

@@ -537,31 +537,12 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         if n_outputs <= 0:
             raise ValueError(f"n_outputs must be positive; got {n_outputs}.")
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
-
-    def _update_init_kwargs(self, **init_kwargs) -> None:
-        """Record constructor arguments that change after construction.
-
-        Helper for :meth:`reset_head` implementations that also change a
-        constructor argument, for example turning a feature extractor into a
-        classifier. Keeps the braindecode init kwargs and the Hugging Face hub
-        config in sync, so :meth:`get_config` and ``save_pretrained`` rebuild
-        the architecture the model actually has.
-
-        Parameters
-        ----------
-        **init_kwargs
-            Constructor argument names mapped to their new values.
-        """
         for config in (
             getattr(self, "_braindecode_init_kwargs", None),
             getattr(self, "_hub_mixin_config", None),
         ):
-            if config is None:
-                continue
-            for name, value in init_kwargs.items():
-                if name in config:
-                    config[name] = value
+            if config is not None and "n_outputs" in config:
+                config["n_outputs"] = n_outputs
 
     def reset_head(self, n_outputs):
         """Replace the classification head for a new number of outputs.
@@ -569,10 +550,6 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=_BraindecodeDocstringMeta):
         This is called automatically by :meth:`from_pretrained` when the
         user passes an ``n_outputs`` that differs from the saved config.
         Override in subclasses that need a model-specific head structure.
-        Implementations keep changed constructor arguments in sync with
-        ``self._update_init_kwargs``, so that a saved model can be loaded
-        back. Implementations requiring positive outputs can also use
-        ``self._set_n_outputs`` to validate and record the new value.
 
         Parameters
         ----------

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/bendr.py

@@ -355,9 +355,7 @@ class BENDR(EEGModuleMixin, nn.Module):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.include_final_layer = True
-        self._update_init_kwargs(final_layer=True)
         self._build_head(n_outputs)
 
     def forward(self, x, return_features=False):

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/biot.py

@@ -237,7 +237,6 @@ class BIOT(EEGModuleMixin, nn.Module):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.final_layer = _ClassificationHead(
             emb_size=self.embed_dim,
             n_outputs=n_outputs,

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/cbramod.py

@@ -224,9 +224,6 @@ class CBraMod(EEGModuleMixin, nn.Module):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
-        # A head implies a classifier, also when built with return_encoder_output.
-        self._update_init_kwargs(return_encoder_output=False)
         if self._n_times is not None and self._n_chans is not None:
             n_patch = self._n_times // self._patch_size
             flat_dim = self._n_chans * n_patch * self._emb_dim

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/eegdino.py

@@ -260,9 +260,7 @@ class EEGDINO(EEGModuleMixin, nn.Module):
         mirroring :meth:`braindecode.models.CBraMod.reset_head`.
         """
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.return_encoder_output = False
-        self._update_init_kwargs(return_encoder_output=False)
         self.final_layer = self._make_head()
 
     def forward(self, x, return_features: bool | None = None):

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/eegpt.py

@@ -371,7 +371,6 @@ class EEGPT(EEGModuleMixin, nn.Module):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.final_layer = _LinearConstraintProbe(
             **self.get_probe_params(),
             n_outputs=n_outputs,

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/labram.py

@@ -823,7 +823,6 @@ class Labram(EEGModuleMixin, nn.Module):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.reset_classifier(n_outputs)
 
     def _adj_temporal_embedding(self, num_ch, batch_size, dim_embed=None):

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/luna.py

@@ -22,7 +22,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
 
-from braindecode.functional import rotate_pairs
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import extract_channel_locations_from_chs_info
 from braindecode.modules.layers import DropPath
@@ -552,9 +551,8 @@ class _RotarySelfAttentionBlock(nn.Module):
         positions = torch.arange(seq_len, device=tensor.device, dtype=tensor.dtype)
         freqs = self.rotary_emb(positions)
 
-        # Keep the full-head reshape (including its empty-input boundary).
         pairs = tensor.reshape(*tensor.shape[:-1], -1, 2)
-        rotated = rotate_pairs(pairs.flatten(-2))
+        rotated = torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).flatten(-2)
         transformed = tensor * freqs.cos() + rotated * freqs.sin()
         return transformed.to(tensor.dtype)
 

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/meta_neuromotor.py

@@ -486,7 +486,6 @@ class MetaNeuromotorHand(EEGModuleMixin, nn.Module, license="cc-by-nc-4.0"):
         """Replace the classification head for a new number of outputs."""
         self.final_layer = nn.Linear(self.final_layer.in_features, n_outputs)
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
 
     def compute_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor:
         """Compute the valid emission length for each input sequence.

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/mvpformer.py

@@ -305,7 +305,6 @@ class MVPFormer(EEGModuleMixin, nn.Module, license="apache-2.0"):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         head_in = (
             self.d_model if self.pooling == "mean" else self.n_chans * self.d_model
         )
@@ -467,11 +466,9 @@ class _MVPAttention(nn.Module):
     # -- shape helpers (segment = "time"/position axis, channel axis) ----------
     @staticmethod
     def _repeat_kv(x, n_rep):
-        # grouped-query attention: replicate each kv head into ``n_rep`` query heads.
-        # Repeat a contiguous copy: on Intel Gaudi (HPU) the einops repeat of the
-        # permuted singleton-axis relative keys returns scrambled values.
+        # grouped-query attention: replicate each kv head into ``n_rep`` query heads
         return repeat(
-            x.contiguous(),
+            x,
             "batch kv_head segment channel head_dim "
             "-> batch (kv_head group) segment channel head_dim",
             group=n_rep,
@@ -508,11 +505,8 @@ class _MVPAttention(nn.Module):
 
     @staticmethod
     def _rel_shift_chan(x):
-        # Relative shift along the channel axis (symmetric distance), as a single
-        # ``torch.gather`` whose backward is a ``scatter_add``. Advanced indexing
-        # would backpropagate through ``index_put_(accumulate=True)``, which Intel
-        # Gaudi (HPU) runs on the host, slowly and with wrong gradients. Index
-        # tensors are built on x.device with long dtype.
+        # Relative shift along the channel axis (symmetric distance). Index
+        # tensors are built on x.device with long dtype for GPU-safe indexing.
         device = x.device
         chan_size = x.shape[-1]
         if chan_size > 1:
@@ -533,9 +527,8 @@ class _MVPAttention(nn.Module):
         shifting_idxes = (chan_size - 1 - shifting_idxes).repeat(
             x.shape[-2] // chan_size, 1
         )
-        return torch.gather(
-            x, -1, shifting_idxes.expand(*x.shape[:-2], *shifting_idxes.shape)
-        )
+        rows = torch.arange(x.size(-2), device=device).unsqueeze(1)
+        return x[..., rows, shifting_idxes]
 
     def _split_heads(self, tensor, num_heads):
         return rearrange(

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/reve.py

@@ -16,7 +16,16 @@ import torch.nn.functional as F
 from einops import rearrange
 from mne.datasets.utils import _get_path
 from torch import nn
-from torch.nn import RMSNorm
+
+# Safe import for older PyTorch versions (Support for Intel-based Macs)
+try:
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    HAS_SDPA = True
+except ImportError:
+    HAS_SDPA = False
+    SDPBackend = None
+    sdpa_kernel = None
 
 from braindecode.models.base import EEGModuleMixin
 
@@ -347,7 +356,6 @@ class REVE(EEGModuleMixin, nn.Module):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self._build_head(n_outputs)
 
     def get_positions(self, channel_names: list[str]) -> torch.Tensor:
@@ -513,11 +521,25 @@ class GEGLU(nn.Module):
         return F.gelu(gates) * x
 
 
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
 class FeedForward(nn.Module):
     def __init__(self, dim: int, hidden_dim: int, geglu: bool):
         super().__init__()
         self.net = nn.Sequential(
-            RMSNorm(dim, eps=1e-6),
+            RMSNorm(dim),
             nn.Linear(dim, hidden_dim * 2 if geglu else hidden_dim, bias=False),
             GEGLU() if geglu else nn.GELU(),
             nn.Linear(hidden_dim, dim, bias=False),
@@ -532,6 +554,59 @@ class FeedForward(nn.Module):
 #################################################################################
 
 
+class ClassicalAttention(nn.Module):
+    def __init__(self, heads: int, use_sdpa: bool | None = None):
+        super().__init__()
+        self.heads = heads
+
+        if use_sdpa is None:
+            self.use_sdpa = HAS_SDPA
+        elif use_sdpa is True and not HAS_SDPA:
+            logger.warning(
+                "SDPA (Scaled Dot Product Attention) was requested, but it is not "
+                "available in your current PyTorch version. Falling back to naive "
+                "implementation. Please upgrade to PyTorch >= 2.2 for SDPA support."
+            )
+            self.use_sdpa = False
+        else:
+            self.use_sdpa = use_sdpa
+
+    def forward(self, qkv: torch.Tensor) -> torch.Tensor:
+        # Split concatenated QKV into separate tensors
+        # qkv shape: (batch, seq_len, 3 * heads * head_dim)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        # Reshape for multi-head attention: split last dim into (heads, head_dim)
+        # (batch, seq_len, heads * head_dim) -> (batch, heads, seq_len, head_dim)
+        q, k, v = (
+            rearrange(
+                t,
+                "batch seq (heads dim) -> batch heads seq dim",
+                heads=self.heads,
+            )
+            for t in (q, k, v)
+        )
+
+        if self.use_sdpa:  # SDPA Implementation
+            with sdpa_kernel(
+                [
+                    SDPBackend.FLASH_ATTENTION,
+                    SDPBackend.EFFICIENT_ATTENTION,
+                    SDPBackend.MATH,
+                ]
+            ):
+                out = F.scaled_dot_product_attention(q, k, v)
+        else:  # Naive Implementation
+            _, _, scale = q.shape[-2], q.device, q.shape[-1] ** -0.5
+            dots = torch.matmul(q, k.transpose(-1, -2)) * scale
+            attn = nn.Softmax(dim=-1)(dots)
+            out = torch.matmul(attn, v)
+
+        # Merge heads back: (batch, heads, seq_len, head_dim) -> (batch, seq_len, heads * head_dim)
+        out = rearrange(out, "batch heads seq dim -> batch seq (heads dim)")
+        return out
+
+
 class Attention(nn.Module):
     """
     Multi-head self-attention layer with RMSNorm.
@@ -541,21 +616,16 @@ class Attention(nn.Module):
         super().__init__()
         inner_dim = head_dim * heads
         self.heads = heads
-        self.norm = RMSNorm(dim, eps=1e-6)
+        self.norm = RMSNorm(dim)
         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
         self.to_out = nn.Linear(inner_dim, dim, bias=False)
 
+        self.attend = ClassicalAttention(self.heads, use_sdpa=True)
+
     def forward(self, x):
         x = self.norm(x)
         qkv = self.to_qkv(x)
-        q, k, v = (
-            rearrange(
-                t, "batch seq (heads dim) -> batch heads seq dim", heads=self.heads
-            )
-            for t in qkv.chunk(3, dim=-1)
-        )
-        out = F.scaled_dot_product_attention(q, k, v)
-        out = rearrange(out, "batch heads seq dim -> batch seq (heads dim)")
+        out = self.attend(qkv)
         return self.to_out(out)
 
 
@@ -633,12 +703,6 @@ class FourierEmb4D(nn.Module):
         self.margin = margin
 
     def forward(self, positions_: torch.Tensor) -> torch.Tensor:
-        # In float32 with autocast off: Intel Gaudi (HPU) autocast also downcasts
-        # the position * frequency products to bf16 before sin/cos (~3 % off).
-        with torch.autocast(device_type=positions_.device.type, enabled=False):
-            return self._embed(positions_.float()).to(positions_.dtype)
-
-    def _embed(self, positions_: torch.Tensor) -> torch.Tensor:
         positions = positions_.clone()
         positions[:, :, -1] *= self.increment_time
         input_shape = positions.shape

{braindecode-1.8.1.dev170227516 → braindecode-1.8.1.dev170248173}/braindecode/models/signal_jepa.py

@@ -643,7 +643,6 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.final_layer = _get_separable_clf_layer(
             conv_layers_spec=self._clf_conv_layers_spec,
             n_chans=self.n_chans,
@@ -888,7 +887,6 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.final_layer = _get_separable_clf_layer(
             conv_layers_spec=self._clf_conv_layers_spec,
             n_chans=self.n_chans,
@@ -1122,7 +1120,6 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
 
     def reset_head(self, n_outputs):
         self._n_outputs = n_outputs
-        self._update_init_kwargs(n_outputs=n_outputs)
         self.final_layer = nn.Sequential(
             nn.Flatten(start_dim=1),
             nn.Linear(self._out_emb_dim, n_outputs),
@@ -1429,14 +1426,8 @@ class _PosEncoder(nn.Module):
             "default_ch_idxs", ch_idxs.to(torch.long), persistent=False
         )
 
-        # Pre-computed tensor for positional encoding on the time dimension, a
-        # non-persistent buffer so that it follows ``.to(device)``:
-        self.encoding_time: torch.Tensor
-        self.register_buffer(
-            "encoding_time",
-            torch.zeros(0, dtype=torch.float32, requires_grad=False),
-            persistent=False,
-        )
+        # Pre-computed tensor for positional encoding on the time dimension:
+        self.encoding_time = torch.zeros(0, dtype=torch.float32, requires_grad=False)
 
     def _check_encoding_time(self, n_times: int):
         if self.encoding_time.size(0) < n_times:
@@ -1474,30 +1465,20 @@ class _PosEncoder(nn.Module):
         assert n_chans_times % n_chans == 0
         n_times = n_chans_times // n_chans
 
-        # Built with torch.cat rather than strided in-place writes into a
-        # new_empty buffer, which Intel Gaudi (HPU) silently dropped.
-        # Channel pos. encoding
-        spat_encoding = self.pos_encoder_spat(ch_idxs)[:, :, None, :].expand(
-            batch_size_chs, n_chans, n_times, self.spat_dim
+        pos_encoding = local_features.new_empty(
+            (batch_size_chs, n_chans, n_times, emb_dim)
         )
+        # Channel pos. encoding
+        pos_encoding[:, :, :, : self.spat_dim] = self.pos_encoder_spat(ch_idxs)[
+            :, :, None, :
+        ]
         # Temporal pos. encoding
         self._check_encoding_time(n_times)
-        time_encoding = self.encoding_time[:n_times].to(
-            device=local_features.device, dtype=local_features.dtype
-        )
-        time_encoding = time_encoding[None, None].expand(
-            batch_size_chs, n_chans, n_times, self.time_dim
+        _ = pos_encoding[:, :, :, self.spat_dim : self.spat_dim + self.time_dim].copy_(
+            self.encoding_time[None, None, :n_times, :],
         )
-        parts = [spat_encoding.to(local_features.dtype), time_encoding]
-        n_rest = emb_dim - self.spat_dim - self.time_dim
-        if n_rest > 0:
-            # Entries new_empty used to leave uninitialised.
-            parts.append(
-                local_features.new_zeros(batch_size_chs, n_chans, n_times, n_rest)
-            )
-        pos_encoding = torch.cat(parts, dim=-1)
 
-        return pos_encoding.reshape(batch_size, n_chans_times, emb_dim)
+        return pos_encoding.view(batch_size, n_chans_times, emb_dim)
 
 
 def _n_times_out(conv_layers_spec, n_times):