braindecode 1.8.1.dev177939635__tar.gz → 1.8.1.dev178237886__tar.gz

{braindecode-1.8.1.dev177939635/braindecode.egg-info → braindecode-1.8.1.dev178237886}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev177939635
+Version: 1.8.1.dev178237886
 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>

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/functional/functions.py

@@ -112,8 +112,11 @@ def prefer_fft_conv(x: torch.Tensor, kernel_size: int) -> bool:
     Rule measured for :class:`~braindecode.models.EEGInceptionMI` (48 to 240
     channels, 128-500 Hz, 2 CPU threads, oneDNN on and off): on CPU the FFT
     wins for float16/bfloat16 inputs (no fast half-precision CPU convolution)
-    and from ``batch_size * kernel_size >= 600``. GPUs and HPUs (no complex
-    dtype) keep the direct convolution.
+    and from ``batch_size * kernel_size >= 600`` if also
+    ``in_channels * kernel_size >= 120``: with one or two input channels a
+    float32 direct convolution stays cheaper up to k = 85 / 31 (13 long
+    temporal convolutions of the model zoo, 22 channels x batch 32). GPUs and
+    HPUs (no complex dtype) keep the direct convolution.
 
     Parameters
     ----------
@@ -123,7 +126,8 @@ def prefer_fft_conv(x: torch.Tensor, kernel_size: int) -> bool:
         Kernel length in samples.
     """
     return x.device.type == "cpu" and (
-        x.element_size() == 2 or x.shape[0] * kernel_size >= 600
+        x.element_size() == 2
+        or (x.shape[0] * kernel_size >= 600 and x.shape[1] * kernel_size >= 120)
     )
 
 

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/models/biot.py

@@ -543,10 +543,15 @@ class _BIOTEncoder(nn.Module):
         emb: Tensor
             (batch_size, emb_size)
         """
+        batch_size, n_chans, n_times = x.shape
+        # One spectrogram for all channels, channel-major so that each
+        # channel's (batch, freq, ts) block is contiguous
+        spec = self.stft(x.transpose(0, 1).reshape(n_chans * batch_size, 1, n_times))
+        spec = spec.unflatten(0, [n_chans, batch_size])
         emb_seq = []
-        for i in range(x.shape[1]):
+        for i in range(n_chans):
             # Getting the spectrogram
-            channel_spec_emb = self.stft(x[:, i : i + 1, :])
+            channel_spec_emb = spec[i]
             # Linear layer to learn some representation over the frequency domain
             # with permutation
             channel_spec_emb = self.patch_embedding(channel_spec_emb)

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/models/codebrain.py

@@ -697,10 +697,35 @@ class _ResidualBlock(nn.Module):
         # Sliding-window attention branch
         # (batch, 2*res_channels, seq_len) -> (batch, seq_len, 2*res_channels)
         h_attn = h_ssm.transpose(1, 2)
-        swa_mask = self.generate_local_window_mask(
-            seq_len, self.swa_window_size, x.device, x.dtype
-        )
-        h_attn, _ = self.attention(h_attn, h_attn, h_attn, attn_mask=swa_mask)
+        if (
+            self.swa_window_size == 1
+            and not self.training
+            and not torch.is_grad_enabled()
+            and h_attn.device.type == "cpu"
+        ):
+            # A one-wide window keeps only the diagonal: the softmax weights are
+            # exactly 1 and 0, so the attention returns its value projection.
+            # Same projection calls as nn.MultiheadAttention (on its (seq,
+            # batch, embed) view), so the output is unchanged bit for bit. (In
+            # training the weights carry dropout; on Gaudi the full attention
+            # was faster.)
+            # ponytail: the unused query/key projections are computed too, to
+            # keep the rounding; in_proj_weight[2E:] alone is ~2x less GEMM.
+            embed_dim = self.attention.embed_dim
+            seq_first = h_attn.transpose(0, 1)
+            value = F.linear(
+                seq_first, self.attention.in_proj_weight, self.attention.in_proj_bias
+            )[..., 2 * embed_dim :]
+            h_attn = (
+                self.attention.out_proj(value.reshape(-1, embed_dim))
+                .view(seq_first.shape)
+                .transpose(0, 1)
+            )
+        else:
+            swa_mask = self.generate_local_window_mask(
+                seq_len, self.swa_window_size, x.device, x.dtype
+            )
+            h_attn, _ = self.attention(h_attn, h_attn, h_attn, attn_mask=swa_mask)
         # (batch, seq_len, 2*res_channels) -> (batch, 2*res_channels, seq_len)
         h_attn = h_attn.transpose(1, 2)
 

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/models/csbrain.py

@@ -710,8 +710,7 @@ class _CSBrainEncoderLayer(nn.Module):
         for start, stop in self.region_bounds:
             region_global = x[:, start:stop, :, :].mean(dim=1, keepdim=True)
             region_global = region_global.permute(0, 2, 1, 3).reshape(batch * T, 1, Fea)
-            for idx in range(start, stop):
-                global_features[:, idx : idx + 1, :] = region_global
+            global_features[:, start:stop, :] = region_global
         x_enhanced = x_flat + self.global_fc(global_features)
 
         attn_output = self.inter_region_attn(

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/models/eegsym.py

@@ -441,6 +441,22 @@ class _AvgPoolTime(nn.Module):
         return y.reshape(b, c, d, y.shape[2], s)
 
 
+class _SpatialConv3d(nn.Conv3d):
+    """``nn.Conv3d`` with a ``(1, 1, k)`` kernel run as ``conv2d`` on the
+    ``(batch, ch, z * time, space)`` view; same parameters. The depthwise
+    ``conv3d`` backward segfaults on Gaudi (synapse 1.21)."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, z, t, s = x.shape
+        y = nn.functional.conv2d(
+            x.reshape(b, c, z * t, s),
+            self.weight.squeeze(2),
+            self.bias,
+            groups=self.groups,
+        )
+        return y.reshape(b, -1, z, t, y.shape[-1])
+
+
 class _InceptionBlock(nn.Module):
     r"""Inception module used in EEGSym architecture.
 
@@ -500,7 +516,7 @@ class _InceptionBlock(nn.Module):
             for _ in range(spatial_resnet_repetitions):
                 self.spatial_convs.append(
                     nn.Sequential(
-                        nn.Conv3d(
+                        _SpatialConv3d(
                             in_channels=filters_per_branch * len(scales_samples),
                             out_channels=filters_per_branch * len(scales_samples),
                             kernel_size=(1, 1, ncha),

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/models/medformer.py

@@ -425,17 +425,13 @@ class _ListPatchEmbedding(nn.Module):
             # Reshape to treat each channel independently: (batch_size * n_chans, 1, n_times)
             x = torch.reshape(x, (batch_size * self.n_chans, 1, self.n_times))
 
+        # add positional embedding to tag each channel (only when not single_channel)
+        if not self.single_channel:
+            x = x + self.channel_embedding(x)
+
         x_list = []
         for padding, value_embedding in zip(self.paddings, self.value_embeddings):
-            x_copy = x.clone()
-            # add positional embedding to tag each channel (only when not single_channel)
-            if not self.single_channel:
-                x_new = x_copy + self.channel_embedding(x_copy)
-            else:
-                x_new = x_copy
-            x_new = padding(x_new).unsqueeze(
-                1
-            )  # (batch_size, 1, enc_in, seq_len+stride)
+            x_new = padding(x).unsqueeze(1)  # (batch_size, 1, enc_in, seq_len+stride)
             x_new = value_embedding(x_new)  # (batch_size, d_model, 1, patch_num)
             x_new = x_new.squeeze(2).transpose(1, 2)  # (batch_size, patch_num, d_model)
             x_list.append(x_new)
@@ -693,16 +689,19 @@ class _EncoderLayer(nn.Module):
         delta: Optional[torch.Tensor] = None,
     ) -> Tuple[List[torch.Tensor], List[Optional[torch.Tensor]]]:
         new_x, attn = self.attention(x, attn_mask=attn_mask, tau=tau, delta=delta)
-        x = [x_orig + self.dropout(x_new) for x_orig, x_new in zip(x, new_x)]
-
-        y = x = [self.norm1(x_val) for x_val in x]
-        y = [
-            self.dropout(self.activation(self.conv1(y_val.transpose(-1, 1))))
-            for y_val in y
-        ]
+        lengths = [x_val.shape[1] for x_val in x]
+        # Norms and kernel-1 convs act per token: run them once on all
+        # granularities. Dropout stays per granularity (same RNG stream).
+        x_all = self.norm1(
+            torch.cat(
+                [x_orig + self.dropout(x_new) for x_orig, x_new in zip(x, new_x)], dim=1
+            )
+        )
+        y = self.activation(self.conv1(x_all.transpose(-1, 1))).split(lengths, dim=-1)
+        y = [self.dropout(y_val) for y_val in y]
         y = [self.dropout(self.conv2(y_val).transpose(-1, 1)) for y_val in y]
 
-        return [self.norm2(x_val + y_val) for x_val, y_val in zip(x, y)], attn
+        return list(self.norm2(x_all + torch.cat(y, dim=1)).split(lengths, dim=1)), attn
 
 
 class _Encoder(nn.Module):

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/models/meta_neuromotor.py

@@ -1075,20 +1075,26 @@ class _RotationInvariantMPFMLP(nn.Module):
             raise ValueError("hidden_dims must be non-empty")
         self.offsets = list(offsets)
         self.activation = activation()
-        self.vectorize = _VectorizeSymmetricMatrix(
+        # Roll both channel axes by each offset, half-vectorize and flatten
+        # (freqs, covariance) as one gather: entry (row, col) of rotation
+        # ``offset`` reads entry ((row - offset) % C, (col - offset) % C).
+        kept = _VectorizeSymmetricMatrix(
             num_channels=num_channels,
             num_adjacent_cov=num_adjacent_cov,
-        )
-        self.flatten_features = Rearrange(
-            "batch time rotation freqs covariance -> "
-            "batch time rotation (freqs covariance)"
-        )
-        # Stacked rotations come out of torch.stack as
-        # (batch, freqs, c1, c2, time, rotation); put time first so the MLP
-        # broadcasts cleanly across (time, rotation).
-        self.rotations_to_time_major = Rearrange(
-            "batch freqs channels1 channels2 time rotation -> "
-            "batch time rotation freqs channels1 channels2"
+        ).flattened_matrix_indices
+        rows, cols = kept // num_channels, kept % num_channels
+        self.register_buffer(
+            "rotation_index",
+            torch.cat(
+                [
+                    freq * num_channels**2
+                    + (rows - offset) % num_channels * num_channels
+                    + (cols - offset) % num_channels
+                    for offset in self.offsets
+                    for freq in range(num_freqs)
+                ]
+            ),
+            persistent=False,
         )
         self.features_to_time_last = Rearrange(
             "batch time features -> batch features time"
@@ -1103,18 +1109,10 @@ class _RotationInvariantMPFMLP(nn.Module):
             dim = hidden_dim
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
-        # Gaudi eager mode rolls a non-contiguous tensor wrongly.
-        inputs = inputs.contiguous()
-        x = torch.stack(
-            [
-                inputs.roll(shifts=[offset, offset], dims=[2, 3])
-                for offset in self.offsets
-            ],
-            dim=-1,
-        )
-        x = self.rotations_to_time_major(x)
-        x = self.vectorize(x)
-        x = self.flatten_features(x)
+        # (batch, freqs, c1, c2, time) -> (batch, time, rotation, features)
+        x = inputs.permute(0, 4, 1, 2, 3).flatten(2)
+        x = x.index_select(-1, self.rotation_index)
+        x = x.unflatten(-1, (len(self.offsets), -1))
         for layer in self.fully_connected_layers:
             x = self.activation(layer(x))
         x = x.mean(dim=2)  # average over the rotations

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/braindecode/modules/dance_modules.py

@@ -10,6 +10,7 @@ import math
 from typing import Dict, Optional
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 from braindecode.functional import sinusoidal_positional_encoding
@@ -172,15 +173,13 @@ class _Attention(nn.Module):
         context = context if context is not None else x
         q = self.to_q(x)
         k, v = self.to_kv(context).chunk(2, dim=-1)
-        # (b, n, h * d) -> (b * h, n, d)
-        q = q.unflatten(-1, (h, -1)).transpose(1, 2).flatten(0, 1)
-        k = k.unflatten(-1, (h, -1)).transpose(1, 2).flatten(0, 1)
-        v = v.unflatten(-1, (h, -1)).transpose(1, 2).flatten(0, 1)
-        attn = (q @ k.transpose(-1, -2) * self.scale).softmax(dim=-1)
-        out = attn @ v
-        # (b * h, n, d) -> (b, n, h * d)
-        out = out.unflatten(0, (-1, h)).transpose(1, 2).flatten(2)
-        return self.to_out(out)
+        # (b, n, h * d) -> (b, h, n, d)
+        q = q.unflatten(-1, (h, -1)).transpose(1, 2)
+        k = k.unflatten(-1, (h, -1)).transpose(1, 2)
+        v = v.unflatten(-1, (h, -1)).transpose(1, 2)
+        out = F.scaled_dot_product_attention(q, k, v, scale=self.scale)
+        # (b, h, n, d) -> (b, n, h * d)
+        return self.to_out(out.transpose(1, 2).flatten(2))
 
 
 class Perceiver(nn.Module):

braindecode-1.8.1.dev178237886/braindecode/version.py

@@ -0,0 +1 @@
+__version__ = "1.8.1.dev178237886"

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886/braindecode.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev177939635
+Version: 1.8.1.dev178237886
 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>

{braindecode-1.8.1.dev177939635 → braindecode-1.8.1.dev178237886}/docs/whats_new.rst

@@ -28,10 +28,17 @@ Current 1.8.1 (2026-08-31)
 Enhancements
 ============
 
+- :class:`braindecode.models.MetaNeuromotorHand` rotates and half-vectorizes the MPF matrices with one precomputed ``index_select`` instead of three rolls, a stack and a permute copy: CPU train step about 2 % faster, Gaudi neutral; outputs, gradients and state dict unchanged bit for bit (:gh:`1292` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.CodeBrain` skips the attention scores in eval mode on CPU at the default ``swa_window_size=1``, where the window keeps only the diagonal and the attention reduces to its value projection: CPU eval forward 1.25x faster at 19 x 6000 (batch 32); training unchanged; outputs unchanged bit for bit (:gh:`1293` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.LUNA` computes its rotary self-attention with :func:`torch.nn.functional.scaled_dot_product_attention`: CPU forward 1-4 % faster, train step neutral to 2.4 % faster; outputs and gradients change by float rounding only, within master's own float32-vs-float64 difference, also with the released ``PulpBio/LUNA`` weights (:gh:`1282` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.DANCE` computes the Perceiver and decoder attention with :func:`torch.nn.functional.scaled_dot_product_attention` instead of an explicit softmax: eval forward 0.7 % faster on CPU and 1.4 % on Gaudi (22 channels, 32 s, batch 32), train step unchanged; outputs and gradients change by float rounding only, within master's own float32-vs-float64 difference (:gh:`1291` by `Bruno Aristimunha`_).
+- :func:`braindecode.functional.prefer_fft_conv` picks the FFT for a float32 convolution only if also ``in_channels * kernel_size >= 120``: with one or two input channels the direct convolution is up to 2.5x cheaper on CPU; :class:`braindecode.models.EEGInceptionMI` decisions and outputs unchanged (:gh:`1283` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.BIOT` computes the spectrograms of all channels with one STFT instead of one per channel: Gaudi forward 1.6x and train step 1.35x faster (22 channels, batch 32), CPU registry forward 3 % faster; outputs and gradients unchanged bit for bit, also with the released weights (:gh:`1284` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.MVPFormer` builds its attention masks and channel-shift index once per forward instead of in every layer: forward 19-25 % faster on Gaudi and 2-12 % on CPU, outputs and gradients unchanged bit for bit (:gh:`1280` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.AttnSleep`, :class:`braindecode.models.Brant`, :class:`braindecode.models.MSVTNet` and :class:`braindecode.models.PBT` no longer compute tensors their forward discards (AttnSleep's first attention product, Brant's reconstruction projection, MSVTNet's branch heads without ``return_features``, PBT's class-token index); outputs and gradients unchanged bit for bit (:gh:`1281` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.CSBrain` copies each region's mean feature into its electrodes with one slice assignment per region instead of one per electrode: train step about 10 % and forward 7 % faster on Gaudi (22 channels, batch 32), unchanged on CPU; outputs are unchanged and gradients change by float rounding only, within master's own float32-vs-float64 difference (:gh:`1285` by `Bruno Aristimunha`_).
 - :class:`braindecode.modules.FilterBankLayer` broadcasts each FIR band filter over the channels instead of repeating it, so its FFT is computed once per band: :class:`braindecode.models.FBCNet`, :class:`braindecode.models.FBMSNet`, :class:`braindecode.models.FBLightConvNet` and :class:`braindecode.models.IFNet` steps 1-10 % faster on CPU and 8-29 % on Gaudi at batch 2, 0-3 % at batch 32; outputs and gradients change by float rounding only at 22 channels, within master's own float32-vs-float64 difference, and not at all at 3 channels (:gh:`1278` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.MEDFormer` adds its channel embedding once instead of once per granularity, and runs its layer norms and kernel-1 feed-forward convolutions once on all granularities (dropout stays per granularity, same masks): train step 3-7 % faster on CPU and 6-7 % on Gaudi; outputs change by float rounding only (:gh:`1288` by `Bruno Aristimunha`_).
 - :func:`braindecode.functional.plv_time` copies the real and imaginary phasor parts once instead of in every matmul for each batch matrix: :class:`braindecode.models.EEGMiner` CPU forward 1.5x and train step 1.4x faster (22 channels, batch 32, 2 threads); outputs and gradients change by float rounding only, within master's own float32-vs-float64 difference (:gh:`1277` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.EEGInceptionMI` runs its max pooling in channels-last layout and pads even kernels once per inception module: CPU forward 1.1-1.25x faster, outputs, gradients and state dict unchanged (:gh:`1267` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.EEGInceptionMI` computes its temporal convolutions with an FFT on CPU when the input is float16/bfloat16 or ``batch_size * largest_kernel >= 600`` (new ``fft_conv`` argument; ``None`` by default, ``False`` keeps the direct convolution bit for bit): float32 CPU steps 1.5-8.8x faster from batch 8 (2 threads), float16/bfloat16 17-117x; outputs change by float rounding (<= 3.6e-7 relative) (:gh:`1268` by `Bruno Aristimunha`_).
@@ -349,6 +356,7 @@ Bug fixes
   when the window length is not divisible by the segment count, instead of
   replacing them with zeros (:gh:`1279` by `Anton Soloviev`_).
 - Model fixes caught by new CPU-only integration checks (complex tensors, host syncs, kernel gaps, device/dtype follow, training after ``inference_mode``, deepcopy/pickle): BrainOmni/BrainTokenizer SELU trains on Gaudi, EEGSym pools with ``avg_pool2d``, FBCNet/FBMSNet/FBLightConvNet and LUNA run in float16, EEGMiner and AttnSleep deep-copy after training, Labram and NeuroRVQ pickle, SignalJEPA heads accept ``channel_strategy``, and tensors built in ``forward`` follow the input in CodeBrain, TCFormer, LUNA, MVPFormer, BrainOmni, ZUNA, DIVER1 and EEGDINO; float32 outputs unchanged (:gh:`1253` by `Bruno Aristimunha`_)
+- :class:`braindecode.models.EEGSym` trains on Gaudi: its depthwise spatial ``Conv3d`` with a ``(1, 1, n_chans)`` kernel runs as a conv2d on a view, since the Gaudi graph compiler segfaulted on the depthwise conv3d backward; same state dict, outputs unchanged, gradients change by float rounding only (:gh:`1289` by `Bruno Aristimunha`_).
 - :func:`braindecode.functional.hilbert_freq` returns float16 for float16 input (it
   returned float32), so :class:`braindecode.models.EEGMiner` runs after
   ``model.to(torch.float16)`` (:gh:`1259` by `Bruno Aristimunha`_)

braindecode-1.8.1.dev177939635/braindecode/version.py

@@ -1 +0,0 @@
-__version__ = "1.8.1.dev177939635"