braindecode 1.8.1.dev174295099__tar.gz → 1.8.1.dev174425547__tar.gz

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

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

@@ -91,16 +91,27 @@ def fft_conv1d(
     n_fft = 2 * _next_fast_len((n_times + kernel_size) // 2)
     xs = spectral_input(x)
     # conv1d is a cross-correlation: a convolution with the flipped kernel
-    x_f = torch.fft.rfft(xs, n=n_fft).permute(2, 0, 1)  # (freq, batch, in)
+    x_f = torch.fft.rfft(xs, n=n_fft)  # (batch, in, freq)
     w_f = torch.fft.rfft(weight.to(xs).flip(-1), n=n_fft).permute(2, 1, 0)
     # Per frequency (batch, in) @ (in, out) in complex numbers, as one real
-    # matmul (a complex one is slower on CPU): [Re x; Im x] @ [Re w, Im w]
-    # (columns interleaved) holds the four real products.
-    p = torch.cat([x_f.real, x_f.imag], 1) @ torch.view_as_real(w_f).flatten(2)
-    p = p.unflatten(1, [2, x.shape[0]]).unflatten(3, [-1, 2])  # x re/im, w re/im
-    y_f = torch.complex(
-        p[:, 0, :, :, 0] - p[:, 1, :, :, 1], p[:, 0, :, :, 1] + p[:, 1, :, :, 0]
-    )
+    # matmul (a complex one is slower on CPU).
+    needs_grad = torch.is_grad_enabled() and (x.requires_grad or weight.requires_grad)
+    if needs_grad or x.shape[0] >= 2 * x.shape[1]:
+        # [Re x, Im x] @ [[Re w, Im w], [-Im w, Re w]] (rows of i * w): the
+        # output columns are the (Re, Im) pairs of the result, read as complex
+        # without a copy (and without a costly backward).
+        w_mat = torch.view_as_real(torch.stack([w_f, w_f * 1j], 1)).flatten(1, 2)
+        y_f = torch.view_as_real(x_f).permute(2, 0, 3, 1).flatten(2) @ w_mat.flatten(2)
+        y_f = torch.view_as_complex(y_f.unflatten(-1, [-1, 2]))
+    else:
+        # Forward only with fewer rows than twice the input channels: the
+        # doubled weight costs more than combining [Re x; Im x] @ [Re w, Im w].
+        x_f = x_f.permute(2, 0, 1)
+        p = torch.cat([x_f.real, x_f.imag], 1) @ torch.view_as_real(w_f).flatten(2)
+        p = p.unflatten(1, [2, x.shape[0]]).unflatten(3, [-1, 2])
+        y_f = torch.complex(
+            p[:, 0, :, :, 0] - p[:, 1, :, :, 1], p[:, 0, :, :, 1] + p[:, 1, :, :, 0]
+        )
     y = torch.fft.irfft(y_f.permute(1, 2, 0), n=n_fft)
     # "same" output t is sample t + kernel_size // 2 of the full convolution
     y = y[..., kernel_size // 2 : kernel_size // 2 + n_times]

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/axon.py

@@ -14,6 +14,7 @@ from torch import nn
 from braindecode.functional import 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
 
 
@@ -462,11 +463,33 @@ class _AxisAttention(nn.Module):
             .unbind(0)
         )
         drop = self.att_drop_prob if self.training else 0.0
-        attended = F.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)
 
+    def _attend_two_scales(
+        self, projected: torch.Tensor, band: torch.Tensor, weights: torch.Tensor
+    ) -> torch.Tensor:
+        """Banded and full attention from one score matrix, mixed per query.
+
+        ``weights`` (N, L, 2) are the scale-gate weights. Both softmaxes share
+        the scores and one product with the values serves the mixture.
+        """
+        n_seq, length, _ = projected.shape
+        query, key, value = (
+            projected.reshape(n_seq, length, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+            .unbind(0)
+        )
+        scores = torch.matmul(query, key.transpose(-2, -1)) * self.head_dim**-0.5
+        weights = weights.reshape(n_seq, 1, length, 2)
+        probs = weights[..., 0:1] * scores.masked_fill(~band, float("-inf")).softmax(
+            -1
+        ) + weights[..., 1:2] * scores.softmax(-1)
+        attended = torch.matmul(probs, value)
+        return attended.transpose(1, 2).reshape(n_seq, length, -1)
+
     def forward(self, x: torch.Tensor, band: Optional[torch.Tensor]) -> torch.Tensor:
         batch, n_chans, n_patches, dim = x.shape
         axis_weights = self.axis_gate(x)  # (B, C, P, 2)
@@ -476,15 +499,26 @@ class _AxisAttention(nn.Module):
         # the restricted branch is the full one; attend once unless the two
         # branches draw their own attention dropout.
         per_channel = self.temporal_qkv(x.reshape(batch * n_chans, n_patches, dim))
-        restricted = self.temporal_out(self._attend(per_channel, band)).reshape(
-            batch, n_chans, n_patches, dim
-        )
-        full = restricted
-        if band is not None or (self.training and self.att_drop_prob > 0):
-            full = self.temporal_out(self._attend(per_channel, None)).reshape(
+        dropout = self.training and self.att_drop_prob > 0
+        if band is not None and not dropout:
+            # The scale-gate weights sum to one, so the two branches share the
+            # scores, the product with the values and temporal_out.
+            attended = self._attend_two_scales(per_channel, band, scale_weights)
+            temporal = self.temporal_out(attended).reshape(
                 batch, n_chans, n_patches, dim
             )
-        temporal = scale_weights[..., 0:1] * restricted + scale_weights[..., 1:2] * full
+        else:
+            restricted = self.temporal_out(self._attend(per_channel, band)).reshape(
+                batch, n_chans, n_patches, dim
+            )
+            full = restricted
+            if dropout:
+                full = self.temporal_out(self._attend(per_channel, None)).reshape(
+                    batch, n_chans, n_patches, dim
+                )
+            temporal = (
+                scale_weights[..., 0:1] * restricted + scale_weights[..., 1:2] * full
+            )
 
         # Spatial path: one sequence of C tokens per time step.
         per_step = self.spatial_qkv(

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/brainmodule.py

@@ -20,6 +20,7 @@ from torch.nn import functional as F
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import has_valid_locations, positions_from_chs_info
 from braindecode.modules import ChannelMerger
+from braindecode.modules.attention import _hpu_lazy
 from braindecode.modules.layers import SubjectLayers
 
 __all__ = ["BrainModule"]
@@ -712,7 +713,7 @@ class _ConvSequence(nn.Module):
                         nn.Conv1d(
                             chout, 2 * chout, 1 + 2 * glu_context, padding=glu_context
                         ),
-                        nn.GLU(dim=1),
+                        _GLU(dim=1),
                     )
                 )
             else:
@@ -727,6 +728,17 @@ class _ConvSequence(nn.Module):
         return x
 
 
+class _GLU(nn.GLU):
+    """``nn.GLU`` that Intel Gaudi's lazy mode can compile."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not torch.jit.is_scripting() and x.device.type == "hpu" and _hpu_lazy():
+            # Its graph compiler fails on ``aten.glu`` after a convolution.
+            a, b = x.chunk(2, dim=self.dim)
+            return a * torch.sigmoid(b)
+        return F.glu(x, self.dim)
+
+
 def _pad_multiple(x: torch.Tensor, base: int) -> torch.Tensor:
     """Pad tensor to be a multiple of base."""
     length = x.shape[-1]

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/eegnex.py

@@ -358,10 +358,14 @@ class EEGNeX(EEGModuleMixin, nn.Module):
 class _TemporalConv(nn.Conv2d):
     """``(1, k)`` Conv2d with "same" output length; on CPU and HPU, through
     :func:`~braindecode.functional.fft_conv1d` where
-    :func:`~braindecode.functional.prefer_fft_conv` expects it to be faster."""
+    :func:`~braindecode.functional.prefer_fft_conv` expects it to be faster.
+    On HPU only when gradients are computed: the FFT pays off in the
+    backward, and Gaudi's direct forward alone is cheaper."""
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if not prefer_fft_conv(x, self.kernel_size[1]):
+        if not prefer_fft_conv(x, self.kernel_size[1]) or (
+            x.device.type == "hpu" and not torch.is_grad_enabled()
+        ):
             return self._conv_forward(x, self.weight, self.bias)
         b, _, h, t = x.shape  # the conv1d of each of the b * h rows
         x = x.transpose(1, 2).reshape(b * h, -1, t)

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/eegpt.py

@@ -13,6 +13,7 @@ from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import DropPath
+from braindecode.modules.attention import _scaled_dot_product_attention
 from braindecode.modules.convolution import Conv1dWithConstraint
 from braindecode.modules.linear import LinearWithConstraint
 from braindecode.util import resolve_montage_name
@@ -949,14 +950,17 @@ class _Attention(nn.Module):
         # 3. Flash Attention
         # Use PyTorch's optimized scaled_dot_product_attention (SDPA)
         # which internally uses FlashAttention or EfficientAttention kernels.
-        y = torch.nn.functional.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=None,
-            dropout_p=self.attn_drop.p if self.training else 0.0,
-            is_causal=self.is_causal,
-        )
+        dropout_p = self.attn_drop.p if self.training else 0.0
+        if self.is_causal:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=True
+            )
+        else:
+            # Gaudi's fused kernel also in lazy mode: faster for EEGPT's
+            # 23-token sequences except on the shortest inputs (256 samples).
+            y = _scaled_dot_product_attention(
+                q, k, v, dropout_p=dropout_p, fused_in_lazy=True
+            )
 
         # Reshape back: (batch, num_heads, seq_len, head_dim) -> (batch, seq_len, embed_dim)
         x = y.transpose(1, 2).flatten(2)

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/medformer.py

@@ -342,17 +342,13 @@ class _PositionalEmbedding(nn.Module):
 
 
 class _CrossChannelTokenEmbedding(nn.Module):
-    def __init__(
-        self, c_in: int, l_patch: int, d_model: int, stride: Optional[int] = None
-    ):
+    def __init__(self, c_in: int, l_patch: int, d_model: int):
         super().__init__()
-        if stride is None:
-            stride = l_patch
         self.token_conv = nn.Conv2d(
             in_channels=1,
             out_channels=d_model,
             kernel_size=(c_in, l_patch),
-            stride=(1, stride),
+            stride=(1, l_patch),
             padding=0,
             padding_mode="circular",
             bias=False,
@@ -364,14 +360,18 @@ class _CrossChannelTokenEmbedding(nn.Module):
                 )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # The kernel spans all c_in rows: run it as a conv1d with c_in input
-        # channels. Gaudi2 rejects conv kernels taller than 256 (c_in = n_times).
-        x = nn.functional.conv1d(
-            x.squeeze(1),
-            self.token_conv.weight.squeeze(1),
-            stride=self.token_conv.stride[1],
+        # The kernel spans all c_in rows and the patches do not overlap, so
+        # the conv is one GEMM of its weight with the patch-reshaped input
+        # (Gaudi2 rejects conv kernels taller than 256 rows, c_in = n_times).
+        l_patch = self.token_conv.kernel_size[1]
+        n_patches = x.shape[-1] // l_patch
+        x = x[:, 0, :, : n_patches * l_patch].reshape(
+            x.shape[0], x.shape[2], n_patches, l_patch
         )
-        return x.unsqueeze(2)
+        x = torch.matmul(
+            x.transpose(1, 2).flatten(2), self.token_conv.weight.flatten(1).t()
+        )
+        return x.transpose(1, 2).unsqueeze(2)
 
 
 class _ListPatchEmbedding(nn.Module):
@@ -493,11 +493,32 @@ class _AttentionLayer(nn.Module):
         _, key_len, _ = keys.shape
         num_heads = self.num_heads
 
-        queries = self.query_projection(queries).view(
-            batch_size, query_len, num_heads, -1
-        )  # multi-head
-        keys = self.key_projection(keys).view(batch_size, key_len, num_heads, -1)
-        values = self.value_projection(values).view(batch_size, key_len, num_heads, -1)
+        if queries is keys and keys is values:
+            # Self-attention: one GEMM for the three projections.
+            weight = torch.cat(
+                [
+                    self.query_projection.weight,
+                    self.key_projection.weight,
+                    self.value_projection.weight,
+                ]
+            )
+            bias = torch.cat(
+                [
+                    self.query_projection.bias,
+                    self.key_projection.bias,
+                    self.value_projection.bias,
+                ]
+            )
+            queries, keys, values = nn.functional.linear(queries, weight, bias).chunk(
+                3, dim=-1
+            )
+        else:
+            queries = self.query_projection(queries)
+            keys = self.key_projection(keys)
+            values = self.value_projection(values)
+        queries = queries.view(batch_size, query_len, num_heads, -1)  # multi-head
+        keys = keys.view(batch_size, key_len, num_heads, -1)
+        values = values.view(batch_size, key_len, num_heads, -1)
 
         out, attn = self.inner_attention(
             queries, keys, values, attn_mask, tau=tau, delta=delta
@@ -561,7 +582,13 @@ class _FullAttention(nn.Module):
         else:
             scale = self.scale
 
-        scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+        # (batch, heads, len, dim) operands: both products are batched GEMMs.
+        queries, keys, values = (
+            queries.transpose(1, 2),
+            keys.transpose(1, 2),
+            values.transpose(1, 2),
+        )
+        scores = torch.matmul(queries, keys.transpose(-2, -1))
 
         if self.mask_flag:
             if attn_mask is None:
@@ -577,7 +604,7 @@ class _FullAttention(nn.Module):
         attention_weights = self.dropout(
             torch.softmax(scale * scores, dim=-1)
         )  # Scaled Dot-Product Attention
-        output_values = torch.einsum("bhls,bshd->blhd", attention_weights, values)
+        output_values = torch.matmul(attention_weights, values).transpose(1, 2)
 
         if self.output_attention:
             return output_values.contiguous(), attention_weights

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/patchedtransformer.py

@@ -12,6 +12,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from braindecode.models.base import EEGModuleMixin
+from braindecode.modules.attention import _scaled_dot_product_attention
 
 
 class PBT(EEGModuleMixin, nn.Module):
@@ -411,13 +412,8 @@ class _MHSA(nn.Module):
         k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
         v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
 
-        y = torch.nn.functional.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=None,
-            dropout_p=self.drop_prob if self.training else 0.0,
-            is_causal=False,
+        y = _scaled_dot_product_attention(
+            q, k, v, dropout_p=self.drop_prob if self.training else 0.0
         )
 
         # re-assemble heads -> (B, T, C)

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/reve.py

@@ -15,10 +15,10 @@ import torch
 import torch.nn.functional as F
 from mne.datasets.utils import _get_path
 from torch import nn
-from torch.nn import RMSNorm
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules.attention import _scaled_dot_product_attention
+from braindecode.modules.layers import RMSNorm
 
 logger = logging.getLogger(__name__)
 
@@ -565,7 +565,8 @@ class Attention(nn.Module):
         q = q.unflatten(-1, (self.heads, -1)).transpose(1, 2)
         k = k.unflatten(-1, (self.heads, -1)).transpose(1, 2)
         v = v.unflatten(-1, (self.heads, -1)).transpose(1, 2)
-        out = _scaled_dot_product_attention(q, k, v)
+        # ~700-token sequences: the fused kernel also beats lazy mode's graph.
+        out = _scaled_dot_product_attention(q, k, v, fused_in_lazy=True)
         return self.to_out(out.transpose(1, 2).flatten(2))
 
 

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/signal_jepa.py

@@ -13,6 +13,10 @@ from einops.layers.torch import Rearrange
 from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
+from braindecode.modules.attention import (
+    _hpu_lazy,
+    _scaled_dot_product_attention,
+)
 
 _DEFAULT_CONV_LAYER_SPEC = (  # downsampling: 128Hz -> 1Hz, receptive field 1.1875s, stride 1s
     (8, 32, 8),
@@ -329,6 +333,56 @@ class _BaseSignalJEPA(EEGModuleMixin, nn.Module):
                 batch_first=True,
             )
 
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training and not torch.jit.is_scripting():
+            return self._encode_train(x)
+        return self.transformer.encoder(x)  # type: ignore
+
+    @torch.jit.unused
+    def _encode_train(self, x: torch.Tensor) -> torch.Tensor:
+        """``self.transformer.encoder(x)`` in training, without SDPA's ``_safe_softmax``.
+
+        With dropout, torch's attention takes SDPA's math path, whose
+        ``_safe_softmax`` adds four passes over the (batch, heads, tokens,
+        tokens) scores to zero the rows a mask hides; there is no mask here.
+        Ops, layouts and dropout draws are otherwise torch's: outputs and
+        gradients are bit-identical on CPU.
+        """
+        encoder = self.transformer.encoder  # type: ignore
+        for layer in encoder.layers:
+            x = layer.norm1(x + layer.dropout1(_self_attention(layer.self_attn, x)))
+            h = layer.dropout(layer.activation(layer.linear1(x)))
+            x = layer.norm2(x + layer.dropout2(layer.linear2(h)))
+        if encoder.norm is not None:
+            x = encoder.norm(x)
+        return x
+
+
+def _self_attention(mha: nn.MultiheadAttention, x: torch.Tensor) -> torch.Tensor:
+    """``mha(x, x, x, need_weights=False)[0]`` (batch first, no mask) as
+    ``F.multi_head_attention_forward`` runs it, with the attention itself
+    through braindecode's SDPA helper (no ``_safe_softmax`` pass)."""
+    x = x.transpose(1, 0)
+    n_tokens, batch_size, emb_dim = x.shape
+    n_heads = mha.num_heads
+    head_dim = emb_dim // n_heads
+    qkv = nn.functional.linear(x, mha.in_proj_weight, mha.in_proj_bias)
+    qkv = qkv.unflatten(-1, (3, emb_dim)).unsqueeze(0).transpose(0, -2)
+    qkv = qkv.squeeze(-2).contiguous()
+    q, k, v = (
+        t.view(n_tokens, batch_size * n_heads, head_dim)
+        .transpose(0, 1)
+        .view(batch_size, n_heads, n_tokens, head_dim)
+        for t in (qkv[0], qkv[1], qkv[2])
+    )
+    out = _scaled_dot_product_attention(
+        q, k, v, dropout_p=mha.dropout if mha.training else 0.0
+    )
+    out = out.permute(2, 0, 1, 3).contiguous()
+    out = out.view(batch_size * n_tokens, emb_dim)
+    out = nn.functional.linear(out, mha.out_proj.weight, mha.out_proj.bias)
+    return out.view(n_tokens, batch_size, emb_dim).transpose(1, 0)
+
 
 class SignalJEPA(_BaseSignalJEPA):
     r"""Architecture introduced in signal-JEPA for self-supervised pre-training, Guetschel, P et al (2024) [1]_
@@ -490,7 +544,7 @@ class SignalJEPA(_BaseSignalJEPA):
         local_features = self.feature_encoder(X)  # type: ignore
         pos_encoding = self.pos_encoder(local_features)  # type: ignore
         local_features += pos_encoding  # type: ignore
-        contextual_features = self.transformer.encoder(local_features)  # type: ignore
+        contextual_features = self._encode(local_features)  # type: ignore
         if return_features:
             out: Dict[str, Optional[torch.Tensor]] = {
                 "features": contextual_features,
@@ -677,7 +731,7 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
         local_features = self.feature_encoder(X)  # type: ignore
         pos_encoding = self.pos_encoder(local_features)  # type: ignore
         local_features += pos_encoding  # type: ignore
-        contextual_features = self.transformer.encoder(local_features)  # type: ignore
+        contextual_features = self._encode(local_features)  # type: ignore
         if return_features:
             out: Dict[str, Optional[torch.Tensor]] = {
                 "features": contextual_features,
@@ -1351,7 +1405,7 @@ class _ConvFeatureEncoder(nn.Sequential):
             return nn.Sequential(
                 conv,
                 nn.Dropout(p=drop_prob),
-                nn.GroupNorm(output_channels, output_channels, affine=True),
+                _GroupNorm(output_channels, output_channels, affine=True),
                 activation(),
             )
         else:
@@ -1361,6 +1415,30 @@ class _ConvFeatureEncoder(nn.Sequential):
         return _n_times_out(self.conv_layers_spec, n_times)
 
 
+class _GroupNorm(nn.GroupNorm):
+    """``nn.GroupNorm``; in HPU eager mode written out with mean and variance.
+
+    In HPU eager mode the backward of torch's group norm compiles new graphs at
+    every call (about 0.2 s per call on the feature encoder's first block);
+    the same normalisation from mean/var ops runs from cached kernels.
+    Elsewhere (CPU, CUDA, HPU lazy) this is ``F.group_norm``.
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not torch.jit.is_scripting() and x.device.type == "hpu" and not _hpu_lazy():
+            xg = x.reshape(x.shape[0], self.num_groups, -1)
+            mean = xg.mean(-1, keepdim=True)
+            var = xg.var(-1, unbiased=False, keepdim=True)
+            x = ((xg - mean) * torch.rsqrt(var + self.eps)).reshape(x.shape)
+            if self.affine:
+                shape = (1, -1) + (1,) * (x.dim() - 2)
+                x = x * self.weight.reshape(shape) + self.bias.reshape(shape)
+            return x
+        return nn.functional.group_norm(
+            x, self.num_groups, self.weight, self.bias, self.eps
+        )
+
+
 class _ChannelEmbedding(nn.Embedding):
     r"""Embedding layer for EEG channels.
 

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/models/zuna.py

@@ -12,12 +12,12 @@ import torch
 from einops import rearrange
 from einops.layers.torch import Rearrange
 from torch import nn
-from torch.nn import RMSNorm
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import extract_channel_locations_from_chs_info
 from braindecode.modules import PatchTokenizer
 from braindecode.modules.attention import _scaled_dot_product_attention
+from braindecode.modules.layers import RMSNorm
 
 
 class ZUNA(EEGModuleMixin, nn.Module, license="apache-2.0"):
@@ -404,18 +404,16 @@ class _RotaryPositionEmbedding(nn.Module):
 
         batch_size, sequence_length, num_heads, head_dim = query.shape
         work = torch.promote_types(query.dtype, torch.float32)
+        # Each (real, imaginary) pair (a, b) rotates to (-b, a): swap the pair
+        # here; the sign is in ``rotary_sine`` (see _ZUNAEncoder).
         query_pairs = query.to(work).reshape(
             batch_size, sequence_length, num_heads, head_dim // 2, 2
         )
         key_pairs = key.to(work).reshape(
             batch_size, sequence_length, num_heads, head_dim // 2, 2
         )
-        rotated_query = self.merge_rotary_pairs(
-            torch.stack((-query_pairs[..., 1], query_pairs[..., 0]), dim=-1)
-        )
-        rotated_key = self.merge_rotary_pairs(
-            torch.stack((-key_pairs[..., 1], key_pairs[..., 0]), dim=-1)
-        )
+        rotated_query = self.merge_rotary_pairs(query_pairs.flip(-1))
+        rotated_key = self.merge_rotary_pairs(key_pairs.flip(-1))
 
         query = (query.to(work) * rotary_cosine + rotated_query * rotary_sine).type_as(
             query
@@ -607,9 +605,12 @@ class _ZUNAEncoder(nn.Module):
             rotary_frequencies.cos(),
             persistent=False,
         )
+        # Sine with the pair rotation's sign folded in: -sin on the real, +sin
+        # on the imaginary component of every pair (exact, saves a negation).
+        pair_sign = torch.tensor([-1.0, 1.0]).repeat(rotary_frequencies.shape[-1] // 2)
         self.register_buffer(
             "rotary_sine",
-            rotary_frequencies.sin(),
+            rotary_frequencies.sin() * pair_sign,
             persistent=False,
         )
 

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/modules/attention.py

@@ -836,31 +836,61 @@ class CATLite(nn.Module):
 
 @torch.jit.unused
 def _hpu_eager_fused_sdpa(
-    query: Tensor, key: Tensor, value: Tensor, scale: Optional[float] = None
+    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.
+    # 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():
+    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
+) -> 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.
+    scale = math.sqrt(1.0 / math.sqrt(query.size(-1)))
+    attn = torch.matmul(query * scale, key.transpose(-2, -1) * scale)
+    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,
+    fused_in_lazy: bool = False,
 ) -> Tensor:
     """``F.scaled_dot_product_attention``; Gaudi's fused kernel on HPU.
 
     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).
+    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 (
         not torch.jit.is_scripting()
@@ -869,9 +899,21 @@ def _scaled_dot_product_attention(
         and dropout_p == 0.0
         and (query.dtype == torch.float32 or query.dtype == torch.bfloat16)
     ):
-        out = _hpu_eager_fused_sdpa(query, key, value)
+        out = _hpu_eager_fused_sdpa(query, key, value, 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)
     return F.scaled_dot_product_attention(
         query, key, value, attn_mask=attn_mask, dropout_p=dropout_p
     )

{braindecode-1.8.1.dev174295099 → braindecode-1.8.1.dev174425547}/braindecode/modules/dance_modules.py

@@ -28,7 +28,19 @@ class _ConvBlock(nn.Module):
         self.act = None if last else activation()
 
     def forward(self, x):
-        x = self.conv(x)
+        # Taps more than (T - 1) // dilation steps off-centre only ever read the
+        # zero padding: drop them (the same sum without its zero terms, fewer
+        # multiply-adds at long dilations).
+        d = self.conv.dilation[0]
+        half = self.conv.kernel_size[0] // 2
+        reach = (x.shape[-1] - 1) // d
+        if reach < half:
+            weight = self.conv.weight[:, :, half - reach : half + reach + 1]
+            x = nn.functional.conv1d(
+                x, weight, self.conv.bias, padding=reach * d, dilation=d
+            )
+        else:
+            x = self.conv(x)
         if self.act is not None:
             x = self.act(x)
         return x