braindecode 1.8.1.dev174425547__tar.gz → 1.8.1.dev174649576__tar.gz

{braindecode-1.8.1.dev174425547/braindecode.egg-info → braindecode-1.8.1.dev174649576}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev174425547
+Version: 1.8.1.dev174649576
 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.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/ctnet.py

@@ -378,7 +378,15 @@ class _PatchEmbeddingEEGNet(nn.Module):
         Tensor
             Embedded patches of shape (batch_size, num_patches, embedding_dim).
         """
-        x = self.eegnet_module(x)
+        layers = self.eegnet_module
+        x = layers[1](layers[0](x))
+        # The depthwise (n_chans, 1) convolution as a matmul over the channel
+        # axis: same values, a faster backward than the grouped convolution's.
+        b, f, c, t = x.shape
+        x = torch.matmul(layers[2].weight.view(f, -1, c), x).view(b, -1, 1, t)
+        for i, layer in enumerate(layers):
+            if i > 2:
+                x = layer(x)
         x = self.projection(x)
         return x
 

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/eegconformer.py

@@ -4,6 +4,7 @@
 import warnings
 
 import torch
+import torch.nn.functional as F
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
@@ -346,7 +347,19 @@ class _PatchEmbedding(nn.Module):
         )
 
     def forward(self, x: Tensor) -> Tensor:
-        x = self.shallownet(x)
+        # Nothing runs between the temporal and the spatial convolution, so
+        # they run as one convolution with the composed weight and bias (as
+        # CombinedConv does), without the (batch, filters, chans, time)
+        # intermediate.
+        conv_time, conv_spat = self.shallownet[0], self.shallownet[1]
+        bias_time, bias_spat = conv_time.bias, conv_spat.bias
+        assert bias_time is not None and bias_spat is not None
+        weight = (conv_time.weight * conv_spat.weight.permute(1, 0, 2, 3)).sum(0)
+        bias = conv_spat.weight.sum((2, 3)) @ bias_time + bias_spat
+        x = F.conv1d(x.squeeze(1), weight, bias).unsqueeze(2)
+        for i, layer in enumerate(self.shallownet):
+            if i > 1:
+                x = layer(x)
         x = self.projection(x)
         return x
 

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/eegsym.py

@@ -457,6 +457,58 @@ class _SpatialConv3d(nn.Conv3d):
         return y.reshape(b, -1, z, t, y.shape[-1])
 
 
+class _BatchNorm3d(nn.BatchNorm3d):
+    """``nn.BatchNorm3d`` with running statistics (as built here); on HPU it
+    normalises the ``(batch, channels, -1)`` view, same statistics and
+    buffers: Gaudi's 5-D batch norm backward is 2-7x slower than its 3-D one
+    on these activations. Elsewhere the same ``batch_norm`` call as torch's."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training and self.num_batches_tracked is not None:
+            self.num_batches_tracked.add_(1)
+        y = nn.functional.batch_norm(
+            x.flatten(2) if x.device.type == "hpu" else x,
+            self.running_mean,
+            self.running_var,
+            self.weight,
+            self.bias,
+            self.training,
+            self.momentum,  # type: ignore[arg-type]
+            self.eps,
+        )
+        return y.view_as(x)
+
+
+class _TemporalConv3d(nn.Conv3d):
+    """``nn.Conv3d`` with a ``(1, k, 1)`` kernel (padding on time only); on HPU
+    it runs as ``conv1d`` on a ``(batch * z * space, channels, time)`` copy of
+    the input, same parameters. Gaudi's conv3d along time, which is not the
+    innermost axis here, is 10-50x slower than that conv1d."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.time_padding = (
+            self.padding if isinstance(self.padding, str) else self.padding[1]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.device.type == "hpu":
+            return self._time_last(x)
+        return self._conv_forward(x, self.weight, self.bias)
+
+    def _time_last(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, z, t, s = x.shape
+        y = nn.functional.conv1d(
+            x.permute(0, 2, 4, 1, 3).reshape(b * z * s, c, t),
+            self.weight.flatten(2),
+            self.bias,
+            padding=self.time_padding,
+            groups=self.groups,
+        )
+        y = y.reshape(b, z, s, y.shape[1], y.shape[2])
+        return y.permute(0, 3, 1, 4, 2).contiguous()
+
+
 class _InceptionBlock(nn.Module):
     r"""Inception module used in EEGSym architecture.
 
@@ -498,13 +550,13 @@ class _InceptionBlock(nn.Module):
         for scale in scales_samples:
             self.temporal_convs.append(
                 nn.Sequential(
-                    nn.Conv3d(
+                    _TemporalConv3d(
                         in_channels=in_channels,
                         out_channels=filters_per_branch,
                         kernel_size=(1, scale, 1),
                         padding="same",
                     ),
-                    nn.BatchNorm3d(filters_per_branch, eps=1e-3, momentum=0.01),
+                    _BatchNorm3d(filters_per_branch, eps=1e-3, momentum=0.01),
                     activation,
                     nn.Dropout(drop_prob),
                 )
@@ -527,7 +579,7 @@ class _InceptionBlock(nn.Module):
                             groups=filters_per_branch * len(scales_samples),
                             bias=False,
                         ),
-                        nn.BatchNorm3d(
+                        _BatchNorm3d(
                             filters_per_branch * len(scales_samples),
                             eps=1e-3,
                             momentum=0.01,
@@ -598,13 +650,13 @@ class _ResidualBlock(nn.Module):
 
         # Temporal convolution
         self.temporal_conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=filters,
                 kernel_size=(1, kernel_size, 1),
                 padding="same",
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -618,7 +670,7 @@ class _ResidualBlock(nn.Module):
                 kernel_size=(1, 1, 1),
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -642,7 +694,7 @@ class _ResidualBlock(nn.Module):
                             groups=1,
                             bias=False,
                         ),
-                        nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+                        _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
                         activation,
                         nn.Dropout(drop_prob),
                     )
@@ -694,14 +746,14 @@ class _TemporalBlock(nn.Module):
         super().__init__()
 
         self.conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=filters,
                 kernel_size=(1, kernel_size, 1),
                 padding="same",
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -762,7 +814,7 @@ class _ChannelMergingBlock(nn.Module):
                         padding=(0, 0, 0),  # Valid padding
                         bias=False,
                     ),
-                    nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+                    _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
                     activation,
                     nn.Dropout(drop_prob),
                 )
@@ -779,7 +831,7 @@ class _ChannelMergingBlock(nn.Module):
                 padding=(0, 0, 0),
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -837,21 +889,21 @@ class _TemporalMergingBlock(nn.Module):
 
         # Residual convolution (collapses time dimension)
         self.residual_conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=in_channels,  # Same channels for residual
                 kernel_size=(1, self.temporal_kernel, 1),  # (Z, Time, Space)
                 padding=(0, 0, 0),  # Valid padding - reduces time to 1
                 bias=False,
             ),
-            nn.BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
 
         # Grouped convolution (collapses time dimension, doubles filters)
         self.grouped_conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=filters,  # Double the channels
                 kernel_size=(1, self.temporal_kernel, 1),  # (Z, Time, Space)
@@ -859,7 +911,7 @@ class _TemporalMergingBlock(nn.Module):
                 padding=(0, 0, 0),
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -908,7 +960,7 @@ class _OutputBlock(nn.Module):
                         padding=(0, 0, 0),
                         bias=False,
                     ),
-                    nn.BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
+                    _BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
                     activation,
                     nn.Dropout(drop_prob),
                 )

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/luna.py

@@ -23,6 +23,7 @@ import torch.nn.functional as F
 
 from braindecode.functional import rotate_pairs, spectral_input
 from braindecode.models.base import EEGModuleMixin
+from braindecode.models.signal_jepa import _GroupNorm
 from braindecode.models.util import extract_channel_locations_from_chs_info
 from braindecode.modules.blocks import PatchTokenizer
 from braindecode.modules.layers import DropPath
@@ -884,7 +885,7 @@ class _PatchEmbedNetwork(nn.Module):
                 stride=(1, self.kernel_size // 2),
                 padding=(0, self.kernel_size // 2 - 1),
             ),
-            nn.GroupNorm(self.groups, self.out_channels),
+            _GroupNorm(self.groups, self.out_channels),
             nn.GELU(),
             nn.Conv2d(
                 in_channels=self.out_channels,
@@ -893,7 +894,7 @@ class _PatchEmbedNetwork(nn.Module):
                 stride=(1, 1),
                 padding=(0, 1),
             ),
-            nn.GroupNorm(self.groups, self.out_channels),
+            _GroupNorm(self.groups, self.out_channels),
             nn.GELU(),
             nn.Conv2d(
                 in_channels=self.out_channels,
@@ -902,7 +903,7 @@ class _PatchEmbedNetwork(nn.Module):
                 stride=(1, 1),
                 padding=(0, 1),
             ),
-            nn.GroupNorm(self.groups, self.out_channels),
+            _GroupNorm(self.groups, self.out_channels),
             nn.GELU(),
         )
 

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/mirepnet.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 from typing import Dict, Optional, Union
 
+import torch.nn.functional as F
 from torch import Tensor, nn
 
 from braindecode.models.base import EEGModuleMixin
@@ -300,15 +301,25 @@ class _PatchEmbedding(nn.Module):
         self.projection = nn.Conv2d(n_filters_spat, embed_dim, kernel_size=(1, 1))
 
     def forward(self, x: Tensor) -> Tensor:
-        x = x.unsqueeze(1)
-        x = self.conv1(x)
-        x = self.conv2(x)
+        # Nothing runs between the temporal and the spatial convolution, so
+        # they run as one convolution over the channels with the composed
+        # weight and bias, without the (batch, filters, chans, time)
+        # intermediate.
+        bias_time, bias_spat = self.conv1.bias, self.conv2.bias
+        assert bias_time is not None and bias_spat is not None
+        w_spat = self.conv2.weight.squeeze(-1).transpose(1, 2)  # (spat, chans, time)
+        weight = w_spat @ self.conv1.weight.flatten(1)  # (spat, chans, kernel)
+        bias = self.conv2.weight.sum((2, 3)) @ bias_time + bias_spat
+        x = F.conv1d(x, weight, bias).unsqueeze(2)
         x = self.bn(x)
         x = self.elu(x)
         x = self.pool(x)
         x = self.dropout(x)
         x = self.projection(x)
-        return x.flatten(start_dim=2).transpose(1, 2)
+        # Contiguous tokens: each residual add takes its first operand's
+        # layout, so a transposed view here would keep every add and layer
+        # norm of the encoder reading strided memory.
+        return x.flatten(start_dim=2).transpose(1, 2).contiguous()
 
 
 class _ResidualAdd(nn.Module):

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/msvtnet.py

@@ -5,6 +5,7 @@ from typing import Tuple, Type, Union
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from einops.layers.torch import Rearrange
 
 from braindecode.models.base import EEGModuleMixin, _too_short_input_error
@@ -286,6 +287,120 @@ class _TSConv(nn.Sequential):
             nn.Dropout(drop_prob),
         )
 
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not torch.is_grad_enabled() and x.device.type == "hpu":
+            # On HPU without autograd the grouped convolution's forward is the
+            # faster one.
+            x = self[2](self[1](self[0](x)))
+        elif x.numel() * self[0].weight.shape[0] >= 2**26:
+            # ponytail: size heuristic. The spatial-first form adds ~50 small ops
+            # per branch; they pay off once the (batch, filters, chans, times)
+            # tensor it avoids reaches 2**26 elements (Gaudi lazy-mode crossover;
+            # eager mode and CPU gain from smaller inputs already).
+            x = self._spatial_first(x)
+        else:
+            # The depthwise (n_channels, 1) convolution as a matmul over the
+            # channel axis: same values, a faster backward than the grouped
+            # convolution's.
+            x = self[1](self[0](x))
+            b, f, c, t = x.shape
+            x = torch.matmul(self[2].weight.view(f, -1, c), x).view(b, -1, 1, t)
+        for i, layer in enumerate(self):
+            if i > 2:
+                x = layer(x)
+        return x
+
+    def _spatial_first(self, x: torch.Tensor) -> torch.Tensor:
+        """Temporal conv, batch norm and depthwise spatial conv (layers 0-2).
+
+        The three are linear per filter, so the spatial weights project the
+        input first and each filter's temporal kernel then runs on its
+        projections; in training the batch norm's statistics come from
+        :func:`_same_conv_moments`. The (batch, filters, chans, times) tensor
+        is never built.
+        """
+        b, _, c, t = x.shape
+        w = self[0].weight.flatten(1)
+        f, k = w.shape
+        w2 = self[2].weight.view(-1, c)
+        d = w2.shape[0] // f
+        if self[1].training or self[1].running_mean is None:
+            mean, var = _same_conv_moments(x.reshape(b * c, t), w)
+            if self[1].training and self[1].track_running_stats:
+                self[1].num_batches_tracked.add_(1)
+                momentum = self[1].momentum
+                if momentum is None:
+                    factor = 1.0 / float(self[1].num_batches_tracked)
+                else:
+                    factor = momentum
+                n = b * c * t
+                self[1].running_mean.mul_(1 - factor).add_(mean.detach(), alpha=factor)
+                self[1].running_var.mul_(1 - factor).add_(
+                    var.detach(), alpha=factor * n / (n - 1)
+                )
+        else:
+            mean, var = self[1].running_mean, self[1].running_var
+        scale = self[1].weight * torch.rsqrt(var + self[1].eps)
+        shift = (self[1].bias - mean * scale).unsqueeze(1) * w2.sum(1).view(f, d)
+        y = F.conv1d(
+            torch.matmul(w2, x.reshape(b, c, t)),
+            w.unsqueeze(1).expand(f, d, k).reshape(f * d, 1, k),
+            padding="same",
+            groups=f * d,
+        )
+        scale = scale.unsqueeze(1).expand(f, d).reshape(-1, 1)
+        return torch.addcmul(shift.view(-1, 1), y, scale).unsqueeze(2)
+
+
+def _same_conv_moments(
+    x: torch.Tensor, w: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Mean and biased variance per filter of ``conv1d(x, w, padding="same")``.
+
+    Over all rows and times, for ``x`` of shape (rows, times) and ``w`` of shape
+    (filters, k), without building the (rows, filters, times) output: the sum of
+    squares is the filter's autocorrelation against the input's (one matmul over
+    blocks of k samples), minus the outputs that "same" crops at both ends. The
+    input is centred first so the variance does not cancel against a large mean.
+    """
+    rows, t = x.shape
+    f, k = w.shape
+    n = rows * t
+    m = x.mean()
+    x = x - m
+    # a[l] = sum over rows and s of x[s] x[s + l], for l < k
+    nb = (t + k - 1) // k
+    blk = F.pad(x, (0, (nb + 1) * k - t)).view(rows, nb + 1, k)
+    nxt = torch.cat([blk[:, :nb], blk[:, 1:, : k - 1]], 2).reshape(-1, 2 * k - 1)
+    a = blk[:, :nb].reshape(-1, k).t() @ nxt
+    a = F.pad(a.reshape(-1), (0, k)).view(k, 2 * k)[:, :k].sum(0)
+    # sum of squares over every output of the full (uncropped) correlation
+    ww = F.conv1d(w.unsqueeze(0), w.unsqueeze(1), padding=k - 1, groups=f)[0]
+    sq = torch.mv(ww, torch.cat([a[1:].flip([0]), a]))
+    if k > 1:
+        right = k - 1 - (k - 1) // 2
+        pad = max(0, k - 1 - t)
+        ends = torch.cat(
+            [
+                F.pad(x[:, : k - 1], (k - 1, pad)),
+                F.pad(x[:, t - min(t, k - 1) :], (pad, k - 1)),
+            ]
+        )
+        ends = F.conv1d(ends.unsqueeze(1), w.unsqueeze(1)).view(2, rows, f, k - 1)
+        ends = ends.square().sum(1)
+        sq = sq - ends[0, :, :right].sum(1) - ends[1, :, right:].sum(1)
+    # undo the centring: conv(x + m) = conv(x) + m * v, v = conv(ones)
+    u = F.conv1d(x.sum(0).view(1, 1, t), w.unsqueeze(1), padding="same")[0]
+    v = F.conv1d(torch.ones_like(x[:1]).unsqueeze(0), w.unsqueeze(1), padding="same")[0]
+    mean_c = u.sum(1) / n
+    dv = v - v.mean(1, keepdim=True)
+    var = (
+        sq / n
+        - mean_c.square()
+        + m * (m * dv.square().mean(1) + 2 * (u * dv).sum(1) / n)
+    )
+    return mean_c + m * v.mean(1), var
+
 
 class _PositionalEncoding(nn.Module):
     r"""

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576}/braindecode/models/neurorvq.py

@@ -16,6 +16,7 @@ from torch import Tensor, nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.labram import _Attention
+from braindecode.models.signal_jepa import _GroupNorm
 from braindecode.modules import MLP, DropPath
 
 # Channel order used to train NeuroRVQ-EEG v1. Keep in sync with the released
@@ -281,7 +282,7 @@ class _MultiScaleTemporalConv(nn.Module):
                 f"conv1_{i}",
                 nn.Conv2d(1, out_chans, (1, kernel), padding=(0, kernel // 2)),
             )
-            setattr(self, f"norm1_{i}", nn.GroupNorm(4, out_chans))
+            setattr(self, f"norm1_{i}", _GroupNorm(4, out_chans))
             setattr(self, f"pool1_{i}", nn.AvgPool2d((1, 2)))
         for i, kernel in enumerate(kernels[1], start=1):
             setattr(
@@ -289,7 +290,7 @@ class _MultiScaleTemporalConv(nn.Module):
                 f"conv2_{i}",
                 nn.Conv2d(out_chans, out_chans, (1, kernel), padding=(0, kernel // 2)),
             )
-            setattr(self, f"norm2_{i}", nn.GroupNorm(4, out_chans))
+            setattr(self, f"norm2_{i}", _GroupNorm(4, out_chans))
             setattr(self, f"pool2_{i}", nn.AvgPool2d((1, 4)))
         self.gelu1 = activation()
         self.gelu2 = activation()

braindecode-1.8.1.dev174649576/braindecode/version.py

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

{braindecode-1.8.1.dev174425547 → braindecode-1.8.1.dev174649576/braindecode.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev174425547
+Version: 1.8.1.dev174649576
 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.dev174425547 → braindecode-1.8.1.dev174649576}/docs/whats_new.rst

@@ -33,12 +33,20 @@ Enhancements
 - :class:`braindecode.models.AXON` computes its banded and full temporal attention from one score matrix: both softmaxes share the query-key product, and one product with the values and one ``temporal_out`` projection serve their scale-gated mixture (the gate weights sum to one). Gaudi train step 1.10x faster in lazy mode and 1.17x in eager mode, CPU 1.05x (21 channels x 6000 samples, batch 64); outputs and gradients differ from before by float rounding of the new summation order, up to 1.6x master's own float32-vs-float64 difference on the outputs (more on the ~1e-9 gate-bias gradients), argmax unchanged; the attention-dropout and no-band paths are unchanged (:gh:`1359` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.EEGInceptionMI` takes the channels-last copy for its max pooling only on CPU: on Intel Gaudi (HPU) the layout copies made the eager train step 2.6x slower than in 1.8.1 (now back to 1.8.1's time); CPU unchanged (:gh:`1344` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.SignalJEPA`, :class:`braindecode.models.SignalJEPA_Contextual`, :class:`braindecode.models.SignalJEPA_PreLocal` and :class:`braindecode.models.SignalJEPA_PostLocal` compute the group norm of their feature encoder from mean and variance in HPU eager mode on Intel Gaudi: torch's group norm costs 120-190 ms per forward and backward there whatever the shape. Gaudi eager train step 0.10x (PreLocal) and 0.17x (PostLocal) at 21 channels x 6000 samples, 0.34x for SignalJEPA at 22 x 1000; same code as before on CPU, CUDA and in HPU lazy mode (:gh:`1364` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.EEGSym` runs its temporal ``(1, k, 1)`` convolutions on Intel Gaudi (HPU) as ``conv1d`` on a time-last copy of the input: Gaudi's conv3d along time, which is not the innermost axis of EEGSym's 5-D activations, is 10-50x slower than that conv1d. Gaudi train step 0.47x in eager and 0.57x in lazy mode at 21 channels x 6000 samples; same code as before on CPU and CUDA, same state dict (:gh:`1375` by `Bruno Aristimunha`_).
 - :func:`braindecode.functional.fft_conv1d` runs on Intel Gaudi (HPU) as overlap-save blocks with the DFT as real matmuls instead of a host FFT, and :func:`braindecode.functional.prefer_fft_conv` picks it there from 8 input channels and ``batch_size * n_times * kernel_size >= 8e6``: Gaudi eager train step of :class:`braindecode.models.EEGInceptionMI` 1.4-2.3x faster and :class:`braindecode.models.EEGNeX` 2.1x faster on 60-s windows at 100 Hz, batch 64; CPU unchanged; Gaudi outputs and gradients change by float rounding only (:gh:`1341` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.EEGNeX` keeps its long temporal convolution direct on Intel Gaudi (HPU) when no gradient is computed: the FFT path pays off only through the backward, so evaluation and inference on Gaudi were 1.8x (eager) and 2.2x (lazy) slower than in 1.8.1 on 60-s windows at 100 Hz, batch 64; now back to 1.8.1's time; training, CPU and GPU unchanged; Gaudi outputs change by float rounding only (:gh:`1363` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.SignalJEPA` and :class:`braindecode.models.SignalJEPA_Contextual` run their transformer encoder in training with the attention through the shared SDPA helper (no ``_safe_softmax`` pass of torch's math path, :gh:`1361`), the layers otherwise op for op as torch's: Gaudi eager train step 0.86x at 21 channels x 6000 samples; bit-identical on CPU, same dropout draws; eval, HPU lazy mode and TorchScript unchanged (:gh:`1365` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.CTNet` computes the depthwise spatial convolution of its patch embedding as a matmul over the channel axis: same forward values on CPU (bit-identical outputs and batch-norm statistics), a faster backward than the grouped-convolution kernels. Train step 0.90x on CPU and 0.92x (eager) / 0.81x (lazy) on Intel Gaudi at 21 channels x 6000 samples, batch 64; 0.85x on CPU at 22 x 1000, batch 32; parameters and state dict unchanged, gradients change by float rounding (:gh:`1374` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.CSBrain` runs its region embedding and the region means of every encoder layer for all regions at once (a 0/1 neighbour-selection matmul and one batched matmul per kernel; one ``index_add_``) instead of one convolution and one mean per region: about half the operators per forward, CPU forward and train step 1.14-1.17x faster, Gaudi forward 1.30-1.45x and train step 1.36-1.56x faster; outputs change by float rounding only (:gh:`1335` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.EEGSym` runs its batch norms on Intel Gaudi (HPU) on the ``(batch, filters, -1)`` view of its 5-D activations: Gaudi's 5-D batch norm backward with the 12 electrodes innermost is 2-7x slower than the 3-D one. Gaudi eager train step 0.75x on top of :gh:`1375` at 21 channels x 6000 samples (lazy 0.95x); same ``batch_norm`` call as before on CPU and CUDA, same state dict (:gh:`1376` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.DIVER1` computes the amplitude spectrum of its spectral embedding on the device on accelerators without complex tensors (Intel Gaudi / HPU) instead of a host FFT round trip: Gaudi forward 1.02-1.12x and train step 1.02-1.10x faster; CPU unchanged (:gh:`1334` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.DIVER1` builds and overlap-adds its STCPE sliding windows with one gather and one ``index_add`` instead of ``unfold``/``fold``: Gaudi forward 1.25-1.36x and train step 1.13-1.32x faster, CPU 0-1 % faster; outputs and released-checkpoint outputs unchanged, gradients change by float rounding only (:gh:`1333` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.EEGConformer` runs the temporal and the spatial convolution of its patch embedding as one convolution over the channels with the composed weight and bias (the two are linear with nothing in between): no (batch, filters, channels, time) intermediate and 3.4x fewer stem FLOPs per train step. Train step 0.58x on CPU and 0.68x (eager) / 0.74x (lazy) on Intel Gaudi at 21 channels x 6000 samples, batch 64; 0.41x on CPU at 22 x 1000, batch 32; parameters and state dict unchanged, outputs change by float rounding (:gh:`1373` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.MIRepNet` runs the temporal and the spatial convolution of its patch embedding as one convolution over the channels with the composed weight and bias (the two are linear with nothing in between): no (batch, filters, channels, time) intermediate and 4.3x fewer stem FLOPs per train step. Train step 0.65x on CPU and 0.63x (eager) / 0.54x (lazy) on Intel Gaudi at 21 channels x 6000 samples, batch 64; 0.60x on CPU at 22 x 1000, batch 32; parameters and state dict unchanged, outputs change by float rounding (:gh:`1377` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.LUNA` computes the three group norms of its patch embedding from mean and variance in HPU eager mode on Intel Gaudi, with the group norm of :class:`braindecode.models.SignalJEPA` (:gh:`1364`): torch's group-norm backward was about half of each eager train step there. Gaudi eager train step 0.53x at 21 channels x 6000 samples (batch 64) and 0.31x at 22 x 1000 (batch 32); CPU, CUDA and Gaudi lazy mode run the same ``group_norm`` call as before (bit-identical on CPU), same state dict (:gh:`1379` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.MIRepNet` makes the tokens of its patch embedding contiguous. The transposed view carried through every residual add of the encoder (a sum takes the memory layout of its first operand), so each add and layer norm read strided memory. Gaudi eager train step 0.92x at 21 channels x 6000 samples, batch 64 (CPU 0.98x, lazy unchanged); outputs change by float rounding (:gh:`1378` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.NeuroRVQ` and :class:`braindecode.models.NeuroRVQTokenizer` compute the group norms of their multi-scale patch embedding from mean and variance in HPU eager mode on Intel Gaudi, as :class:`braindecode.models.SignalJEPA` does (:gh:`1364`): the backward of torch's group norm compiled new graphs at every call there (8 per train step, about 70 ms each). Gaudi eager train step 0.22x (NeuroRVQ) and 0.39x (NeuroRVQTokenizer) at 21 channels x 800 samples, batch 64; same code as before on CPU, CUDA and in HPU lazy mode (:gh:`1382` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.REVE` computes the 4-D position embedding of its default channel positions once per forward and broadcasts it over the batch instead of computing it for one identical copy per example: 3 ms instead of 10-12 ms of each Gaudi train step (21 channels x 6000 samples, batch 64; step 1.01x); explicit ``pos`` inputs are unchanged (:gh:`1347` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.ContraWR` computes its STFT magnitude on the device on accelerators without complex tensors (Intel Gaudi / HPU) instead of a host STFT and a host-to-device copy of the spectrogram: Gaudi forward 4.2-7.1x and train step 2.4-5.2x faster; CPU unchanged (:gh:`1332` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.NeuroRVQTokenizer` sums the per-code counts and vectors of its EMA codebook update with ``index_add_`` instead of a dense one-hot matrix and GEMM: CPU train step 1.3x faster (22 channels, batch 32, 2 threads), Gaudi unchanged; outputs, gradients and token ids unchanged, the updated codebooks change by float rounding only (:gh:`1308` by `Bruno Aristimunha`_).
@@ -68,6 +76,8 @@ Enhancements
 - :class:`braindecode.models.DANCE` skips the taps of its dilated convolutions that only ever read the zero padding (kernel offset at or beyond the input length): CPU train step 0.92x and Gaudi eager 0.97x (21 channels x 6000 samples, batch 64); outputs change by float rounding only (:gh:`1343` by `Bruno Aristimunha`_).
 - :class:`braindecode.modules.ChannelMerger` (DANCE, ``BrainModule(use_merger=True)``) computes its channel mix as ``(x^T @ weights^T)^T`` on HPU, where Gaudi's lazy-mode graph compiler failed on ``weights @ x`` feeding a convolution's backward: DANCE now trains in Gaudi lazy mode (train step 622 ms vs 844 ms in eager mode, 21 channels x 6000 samples, batch 64); other devices run the same code as before (:gh:`1351` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.BrainModule` computes its GLU gates as ``a * sigmoid(b)`` in Intel Gaudi (HPU) lazy mode, where the graph compiler failed on ``aten.glu`` after the gate's convolution: BrainModule now trains in Gaudi lazy mode (train step 306 ms vs 554 ms in eager mode, 21 channels x 6000 samples, batch 64); eager mode and other devices run the same code as before (:gh:`1371` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.MSVTNet` computes the depthwise spatial convolution of each branch as a matmul over the channel axis when gradients are computed (and on CPU): same values, a faster backward than the grouped convolution's. Gaudi train step 0.86x in eager and 0.87x in lazy mode, CPU 0.88x at 22 channels x 1000 samples (batch 32) and 0.98x at 21 x 6000 (batch 64); evaluation on Gaudi keeps the grouped convolution; outputs and batch-norm statistics unchanged on CPU, gradients change by float rounding (:gh:`1380` by `Bruno Aristimunha`_).
+- :class:`braindecode.models.MSVTNet` with large inputs (when a branch's (batch, filters, channels, times) activation would reach 2**26 elements, e.g. 21 channels x 6000 samples at batch 64) applies each branch's spatial weights to the input first and runs each filter's temporal kernel on its projections, with the first batch norm's statistics computed from the input's autocorrelation, so that activation is never built: same function and state dict, numerical differences only. Train step at 21 x 6000, batch 64: Intel Gaudi 0.57x in eager mode and 0.78x in lazy mode, CPU 0.23x (:gh:`1384` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.Labram` and :class:`braindecode.models.BaRISTA` run attention through the Gaudi bridge's fused SDPA kernel in HPU eager mode (without a mask or attention dropout; lazy mode, CPU and other devices unchanged): Gaudi eager train step Labram 8 % and BaRISTA 5 % faster (21 channels x 6000 samples), differences within the Gaudi float32 error (:gh:`1342` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.DANCE` runs its Perceiver attention through the same Gaudi fused SDPA path in HPU eager mode (torch's SDPA there computes the full score matrix op by op, which made the DANCE Gaudi eager train step 8 % slower than in 1.8.1): Gaudi eager train step 848 -> 763 ms (21 channels x 6000 samples, batch 64); CPU and Gaudi lazy mode unchanged (:gh:`1340` by `Bruno Aristimunha`_).
 - :class:`braindecode.models.REVE` runs its attention through the same helper, so HPU eager mode uses the Gaudi bridge's fused SDPA kernel instead of torch's op-by-op SDPA through the full score matrix (lazy mode, CPU and other devices unchanged): Gaudi eager train step 1.42x faster (21 channels x 6000 samples, batch 64); differences within the Gaudi float32 error (:gh:`1345` by `Bruno Aristimunha`_).

braindecode-1.8.1.dev174425547/braindecode/version.py

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