Braindecode 1.8.1.dev180761562__tar.gz → 1.8.1.dev180835552__tar.gz

{braindecode-1.8.1.dev180761562/braindecode.egg-info → braindecode-1.8.1.dev180835552}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev180761562
+Version: 1.8.1.dev180835552
 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.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/augmentation/__init__.py

@@ -14,6 +14,7 @@ from .transforms import (
     FrequencyShift,
     FTSurrogate,
     GaussianNoise,
+    LabelSmoothing,
     MaskEncoding,
     Mixup,
     SegmentationReconstruction,
@@ -40,6 +41,7 @@ __all__ = [
     "CutMix",
     "ChannelsDropout",
     "GaussianNoise",
+    "LabelSmoothing",
     "ChannelsSymmetry",
     "SmoothTimeMask",
     "BandstopFilter",

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/augmentation/functional.py

@@ -19,7 +19,7 @@ from mne.filter import notch_filter
 from scipy.interpolate import Rbf
 from sklearn.utils import check_random_state
 from torch.fft import fft, ifft
-from torch.nn.functional import pad
+from torch.nn.functional import one_hot, pad
 
 
 def _max_float_dtype(device: torch.device | str) -> torch.dtype:
@@ -1089,10 +1089,12 @@ def mixup(
         Online: https://arxiv.org/abs/1710.09412
     .. [2] https://github.com/facebookresearch/mixup-cifar10/blob/master/train.py
     """
-    device = X.device
     batch_size, n_channels, n_times = X.shape
 
-    X_mix = torch.zeros((batch_size, n_channels, n_times)).to(device)
+    X_mix = X.new_zeros(
+        (batch_size, n_channels, n_times),
+        dtype=X.dtype if X.is_floating_point() else torch.get_default_dtype(),
+    )
     y_a = y.clone()
     y_b = y[idx_perm].clone()
 
@@ -1150,6 +1152,97 @@ def cutmix(
     return X_mix, (y.clone(), y[idx_perm].clone(), lam)
 
 
+def label_smoothing(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    epsilon: float,
+    num_classes: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Smooth targets toward the uniform distribution over the classes.
+
+    See [1]_ for details.
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EEG data in form ``batch_size, n_channels, n_times``. Returned
+        unchanged.
+    y : torch.Tensor
+        Target of length ``batch_size`` holding class indices (float values
+        are accepted when they are integral), or a floating-point tensor of
+        shape ``batch_size, num_classes`` holding target distributions.
+    epsilon : float
+        Smoothing parameter between 0 and 1: the weight of the uniform
+        distribution mixed into the targets. ``0`` returns the unsmoothed
+        targets.
+    num_classes : int | None, optional
+        Number of classes. Required when ``y`` holds class indices; when
+        ``y`` already holds distributions, it must match the size of the
+        last dimension of ``y``. Defaults to None.
+
+    Returns
+    -------
+    tuple
+        ``X``, ``y``. Where ``X`` is unchanged and ``y`` is a
+        floating-point tensor of shape ``batch_size, num_classes`` whose
+        rows sum to one.
+
+    References
+    ----------
+    .. [1] Christian Szegedy, Vincent Vanhoucke, Jonathon Shlens and
+        Zbigniew Wojna (2016). Rethinking the Inception Architecture for
+        Computer Vision. In 2016 IEEE Conference on Computer Vision and
+        Pattern Recognition (CVPR), pp. 2818-2826.
+        Online: https://arxiv.org/abs/1512.00567
+    """
+    if not isinstance(epsilon, Real) or not 0.0 <= float(epsilon) <= 1.0:
+        raise ValueError(
+            f"epsilon should be a real number between 0 and 1. Got {epsilon!r}."
+        )
+    if num_classes is not None and (
+        not isinstance(num_classes, int)
+        or isinstance(num_classes, bool)
+        or num_classes < 2
+    ):
+        raise ValueError(
+            f"num_classes should be an integer of at least 2. Got {num_classes!r}."
+        )
+
+    if y.dim() == 1:
+        # class-index targets: smooth their one-hot encoding
+        if num_classes is None:
+            raise ValueError("num_classes is required when y holds class indices.")
+        if not torch.all(y == y.floor()):
+            raise ValueError(
+                "y holds non-integral values; class-index targets should be "
+                "integers or integral floats."
+            )
+        if y.numel() > 0 and (y.min() < 0 or y.max() >= num_classes):
+            raise ValueError(f"y holds class indices outside [0, {num_classes}).")
+        one_hot_y = one_hot(y.long(), num_classes).to(dtype=X.dtype)
+        return X, one_hot_y * (1.0 - epsilon) + epsilon / num_classes
+
+    if y.is_floating_point() and y.dim() >= 2:
+        # target distributions: interpolate them toward the uniform one
+        n_classes = y.shape[-1]
+        if num_classes is not None and num_classes != n_classes:
+            raise ValueError(
+                f"num_classes ({num_classes}) does not match the size of the "
+                f"last dimension of y ({n_classes})."
+            )
+        if n_classes < 2:
+            raise ValueError(
+                f"y should hold distributions over at least 2 classes. Got {n_classes}."
+            )
+        return X, y * (1.0 - epsilon) + epsilon / n_classes
+
+    raise ValueError(
+        "y should hold class indices or floating-point target distributions "
+        f"with a class dimension; got dtype {y.dtype} and shape "
+        f"{tuple(y.shape)}."
+    )
+
+
 def segmentation_reconstruction(
     X: torch.Tensor,
     y: torch.Tensor,

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/augmentation/transforms.py

@@ -30,6 +30,7 @@ from .functional import (
     frequency_shift,
     ft_surrogate,
     gaussian_noise,
+    label_smoothing,
     mask_encoding,
     mixup,
     segmentation_reconstruction,
@@ -1155,6 +1156,98 @@ class CutMix(Mixup):
         return dict(params, cut_t_start=start, cut_t_len=end - start)
 
 
+class LabelSmoothing(Transform):
+    """Label smoothing from Szegedy et al. (2016) [1]_.
+
+    Replaces each target with its convex combination with the uniform
+    distribution over the classes, penalizing over-confident predictions.
+    The signals are returned unchanged; only the targets change.
+
+    The smoothed targets are floating-point distributions of shape
+    ``batch_size, num_classes`` (each row sums to one), so the training
+    loss must accept target distributions, e.g. soft cross-entropy
+    ``-(y * torch.log_softmax(logits, dim=-1)).sum(-1).mean()`` or
+    ``torch.nn.KLDivLoss``. The transform composes before target-mixing
+    transforms such as :class:`braindecode.augmentation.Mixup`, whose
+    fractional-target path then mixes the smoothed distributions. The
+    fixed-expansion mode (``n_augmentation > 0``) of
+    :class:`braindecode.augmentation.AugmentedDataLoader` is not supported:
+    it cannot concatenate the unsmoothed original targets with the reshaped
+    smoothed copies.
+
+    Parameters
+    ----------
+    epsilon : float
+        Smoothing parameter between 0 and 1: the weight of the uniform
+        distribution mixed into the targets. ``0`` disables smoothing.
+    num_classes : int | None, optional
+        Number of classes. Required when the targets are class indices;
+        when the targets are already distributions of shape
+        ``batch_size, num_classes``, it must match their last dimension.
+        Defaults to None.
+    random_state : int | numpy.random.Generator, optional
+        Seed to be used to instantiate numpy random number generator
+        instance. The transform is deterministic, so the seed only feeds
+        the application-probability machinery, which is inactive here.
+        Defaults to None.
+
+    References
+    ----------
+    .. [1] Christian Szegedy, Vincent Vanhoucke, Jonathon Shlens and
+       Zbigniew Wojna (2016). Rethinking the Inception Architecture for
+       Computer Vision. In 2016 IEEE Conference on Computer Vision and
+       Pattern Recognition (CVPR), pp. 2818-2826.
+       Online: https://arxiv.org/abs/1512.00567
+    """
+
+    operation = staticmethod(label_smoothing)  # type: ignore[assignment]
+
+    def __init__(self, epsilon, num_classes=None, random_state=None):
+        super().__init__(
+            probability=1.0,  # LabelSmoothing has to be applied to whole batches
+            random_state=random_state,
+        )
+        if not isinstance(epsilon, Real) or not 0.0 <= float(epsilon) <= 1.0:
+            raise ValueError(
+                f"epsilon should be a real number between 0 and 1. Got {epsilon!r}."
+            )
+        if num_classes is not None and (
+            not isinstance(num_classes, int)
+            or isinstance(num_classes, bool)
+            or num_classes < 2
+        ):
+            raise ValueError(
+                f"num_classes should be an integer of at least 2. Got {num_classes!r}."
+            )
+        self.epsilon = epsilon
+        self.num_classes = num_classes
+
+    def get_augmentation_params(self, *batch):
+        """Return transform parameters.
+
+        Parameters
+        ----------
+        X : tensor.Tensor
+            The data.
+        y : tensor.Tensor
+            The labels.
+
+        Returns
+        -------
+        params : dict
+            Contains
+
+            * epsilon : float
+                Smoothing parameter between 0 and 1.
+            * num_classes : int | None
+                The number of classes.
+        """
+        return {
+            "epsilon": self.epsilon,
+            "num_classes": self.num_classes,
+        }
+
+
 class SegmentationReconstruction(Transform):
     """Segmentation Reconstruction from Lotte (2015) [Lotte2015]_.
 

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/functional/initialization.py

@@ -1,25 +1,30 @@
 import math
 
+import torch
 from torch import nn
 
 
 def glorot_weight_zero_bias(model):
-    """Initialize parameters of all modules by initializing weights with.
+    """Set the weights of BatchNorm layers to one and all biases to zero.
 
-    glorot uniform/xavier initialization, and setting biases to zero. Weights from
-    batch norm layers are set to 1.
+    Despite its name, this function does not apply the Glorot (Xavier)
+    initialization: other weights keep the initialization of their layer
+    (PyTorch's default, or Xavier uniform for the max-norm constrained layers
+    such as :class:`braindecode.modules.Conv2dWithConstraint`). The name is
+    kept for backward compatibility.
 
     Parameters
     ----------
     model : Module
     """
     for module in model.modules():
-        if hasattr(module, "weight"):
-            if "BatchNorm" in module.__class__.__name__:
-                nn.init.constant_(module.weight, 1)
-        if hasattr(module, "bias"):
-            if module.bias is not None:
-                nn.init.constant_(module.bias, 0)
+        # ``weight``/``bias`` can be None (e.g. BatchNorm with affine=False)
+        weight = getattr(module, "weight", None)
+        if isinstance(weight, torch.Tensor) and "BatchNorm" in type(module).__name__:
+            nn.init.constant_(weight, 1)
+        bias = getattr(module, "bias", None)
+        if isinstance(bias, torch.Tensor):
+            nn.init.constant_(bias, 0)
 
 
 def rescale_parameter(param, layer_id):

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/atcnet.py

@@ -444,8 +444,10 @@ class ATCNet(EEGModuleMixin, nn.Module):
         # ----- Sliding window -----
         sw_concat: list[torch.Tensor] = []  # to store sliding window outputs
         # for w in range(self.n_windows):
-        for idx, (attention, tcn_module, final_layer) in enumerate(
-            zip(self.attention_blocks, self.temporal_conv_nets, self.final_layer)
+        # (final_layer is not zipped in: with concat=True it holds a single
+        # layer, which would stop the loop after the first window)
+        for idx, (attention, tcn_module) in enumerate(
+            zip(self.attention_blocks, self.temporal_conv_nets)
         ):
             conv_feat_w = conv_feat[..., idx : idx + self.Tw]
             # Dimension: (batch_size, F2, Tw)
@@ -458,14 +460,15 @@ class ATCNet(EEGModuleMixin, nn.Module):
             tcn_feat = tcn_module(att_feat)[..., -1]
             # Dimension: (batch_size, F2)
 
-            # Outputs of sliding window can be either averaged after being
-            # mapped by dense layer or concatenated then mapped by a dense
-            # layer
-            if not self.concat:
-                tcn_feat = final_layer(tcn_feat)
-
             sw_concat.append(tcn_feat)
 
+        # Outputs of sliding window can be either averaged after being
+        # mapped by dense layer or concatenated then mapped by a dense
+        # layer
+        if not self.concat:
+            for idx, final_layer in enumerate(self.final_layer):
+                sw_concat[idx] = final_layer(sw_concat[idx])
+
         # ----- Aggregation and prediction -----
         if self.concat:
             sw_concat_agg = torch.cat(sw_concat, dim=1)

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/biot.py

@@ -186,10 +186,12 @@ class BIOT(EEGModuleMixin, nn.Module):
             warn(
                 "The hop length is larger than the sampling frequency. "
                 + "This may cause aliasing. Consider using a smaller "
-                "hop length.",
+                "hop length. Using sfreq // 2 instead.",
                 UserWarning,
             )
-            hop_length = self.sfreq // 2
+            # torch.stft needs an int, and sfreq is usually a float
+            hop_length = int(self.sfreq // 2)
+            self.hop_length = hop_length
 
         if self.input_window_seconds < 1.0:
             warning_msg = (

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/ctnet.py

@@ -181,6 +181,7 @@ class CTNet(EEGModuleMixin, nn.Module, license="mit"):
             drop_prob=self.att_positional_drop_prob,
             n_times=self.n_times,
             pool_size=self.pool_size_1,
+            sequence_length=self.sequence_length,
         )
 
         self.trans = _TransformerEncoder(
@@ -220,7 +221,7 @@ class CTNet(EEGModuleMixin, nn.Module, license="mit"):
         cnn = self.position(cnn)
         trans = self.trans(cnn)
         features = cnn + trans
-        flatten_feature = self.flatten(features)
+        flatten_feature = self.flatten_drop_layer(features)
         out = self.final_layer(flatten_feature)
         return out
 
@@ -516,18 +517,23 @@ class _PositionalEncoding(nn.Module):
         length: int = 100,
         drop_prob: float = 0.1,
         pool_size: int = 8,
+        sequence_length: int | None = None,
     ):
         super().__init__()
         self.pool_size = pool_size
         self.n_times = n_times
 
-        if int(n_times / (pool_size * pool_size)) > length:
+        # n_times / pool_size**2 is the sequence length only when both pooling
+        # sizes are equal; also cover the actual length (e.g. pool_size_2 <
+        # pool_size_1), so that the encoding is never shorter than the input
+        needed = max(int(n_times / (pool_size * pool_size)), sequence_length or 0)
+        if needed > length:
             warn(
                 "the temporal dimensional is too long for this default length. "
                 "The length parameter will be automatically adjusted to "
                 "avoid inference issues."
             )
-            length = int(n_times / (pool_size * pool_size))
+            length = needed
 
         self.dropout = nn.Dropout(drop_prob)
         self.encoding = nn.Parameter(torch.randn(1, length, emb_size))

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/labram.py

@@ -296,7 +296,9 @@ class Labram(EEGModuleMixin, nn.Module):
     learned_patcher : bool
         Whether to use a learned patch embedding (via a convolutional layer) or a fixed patch embedding (via rearrangement).
     embed_dim : int
-        The dimension of the embedding.
+        The dimension of the embedding. In the Neural Tokenizer mode, it is
+        set by the temporal convolution (``conv_out_channels`` and
+        ``patch_size``), with a warning if the given value differs.
     conv_in_channels : int
         The number of convolutional input channels.
     conv_out_channels : int
@@ -515,6 +517,14 @@ class Labram(EEGModuleMixin, nn.Module):
             out = self.patch_embed(dummy)
         # out.shape for tokenizer: (1, n_chans, emb_dim)
         # for decoder:        (1, n_patch, patch_size, emb_dim), but we want last dim
+        if neural_tokenizer and out.shape[-1] != self.embed_dim:
+            warn(
+                f"In the Neural Tokenizer mode the embedding dimension is set "
+                f"by the temporal convolution (conv_out_channels, patch_size): "
+                f"using embed_dim={out.shape[-1]} instead of the given "
+                f"embed_dim={self.embed_dim}.",
+                UserWarning,
+            )
         self.embed_dim = out.shape[-1]
         self.num_features = self.embed_dim
 

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/luna.py

@@ -113,6 +113,9 @@ class LUNA(EEGModuleMixin, nn.Module, license="apache-2.0"):
         Normalization layer class. Default: nn.LayerNorm.
     drop_path : float
         Stochastic depth rate. Default: 0.0.
+    attn_drop : float
+        Dropout rate on the attention weights of the transformer encoder
+        blocks (the original LUNA uses 0). Default: 0.0.
 
     References
     ----------
@@ -222,7 +225,7 @@ class LUNA(EEGModuleMixin, nn.Module, license="apache-2.0"):
                     mlp_ratio=self.mlp_ratio,
                     qkv_bias=True,
                     drop=0.0,
-                    attn_drop=0.0,
+                    attn_drop=self.attn_drop,
                     drop_path=self.drop_path,
                     norm_layer=norm_layer,
                 )

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/medformer.py

@@ -235,17 +235,18 @@ class MEDFormer(EEGModuleMixin, nn.Module, license="mit"):
         self.patch_len_list = patch_len_list
         stride_list = patch_len_list  # Using the same values for strides.
         self.stride_list = stride_list
+        # patches run over time (one token per patch of samples)
         patch_num_list = [
-            int((self.n_chans - patch_len) / stride + 2)
+            int((self.n_times - patch_len) / stride + 2)
             for patch_len, stride in zip(patch_len_list, stride_list)
         ]
         self.patch_num_list = patch_num_list
 
         # Initialize the embedding layer.
         self.enc_embedding = _ListPatchEmbedding(
-            enc_in=self.n_times,
+            enc_in=self.n_chans,
             d_model=self.embed_dim,
-            seq_len=self.n_chans,
+            seq_len=self.n_times,
             patch_len_list=self.patch_len_list,
             stride_list=self.stride_list,
             dropout=self.drop_prob,
@@ -362,7 +363,7 @@ class _CrossChannelTokenEmbedding(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # 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).
+        # (Gaudi2 rejects conv kernels taller than 256 rows).
         l_patch = self.token_conv.kernel_size[1]
         n_patches = x.shape[-1] // l_patch
         x = x[:, 0, :, : n_patches * l_patch].reshape(
@@ -411,10 +412,9 @@ class _ListPatchEmbedding(nn.Module):
         ]
         self.value_embeddings = nn.ModuleList(linear_layers)
         self.position_embedding = _PositionalEmbedding(d_model=d_model)
-        # One row per time sample: windows longer than 5000 samples need more.
-        self.channel_embedding = _PositionalEmbedding(
-            d_model=seq_len, max_len=max(5000, enc_in)
-        )
+        # One sinusoidal code along time per channel, which tags each channel
+        # (the rows past enc_in that the original max_len=5000 adds are unused)
+        self.channel_embedding = _PositionalEmbedding(d_model=seq_len, max_len=enc_in)
         self.dropout = nn.Dropout(dropout)
 
         self.learnable_embeddings = nn.ParameterList(
@@ -423,15 +423,12 @@ class _ListPatchEmbedding(nn.Module):
 
     def forward(
         self, x: torch.Tensor
-    ) -> List[torch.Tensor]:  # (batch_size, seq_len, enc_in)
-        x = x.permute(0, 2, 1)  # (batch_size, enc_in, seq_len)
+    ) -> List[torch.Tensor]:  # (batch_size, enc_in, seq_len) = (B, n_chans, n_times)
+        # The original code takes (batch_size, seq_len, enc_in) and permutes it
+        # to this layout; braindecode inputs already have it.
         if self.single_channel:
-            # After permute: x.shape = (batch_size, n_times, n_chans)
-            # We want to process each channel independently
+            # Treat each channel independently: (batch_size * n_chans, 1, n_times)
             batch_size = x.shape[0]
-            # Permute to get channels in the middle: (batch_size, n_chans, n_times)
-            x = x.permute(0, 2, 1)
-            # 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)

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/sinc_shallow.py

@@ -123,6 +123,12 @@ class SincShallowNet(EEGModuleMixin, nn.Module):
         if activation is None:
             activation = nn.ELU()
 
+        if self.n_times is None:
+            raise ValueError(
+                "SincShallowNet needs to know the input length: pass n_times "
+                "(or input_window_seconds and sfreq)."
+            )
+
         # Define low frequencies for the SincFilter
         low_freqs = torch.arange(
             first_freq,
@@ -139,6 +145,15 @@ class SincShallowNet(EEGModuleMixin, nn.Module):
         else:
             raise ValueError("Padding must be 'valid' or 'same'.")
 
+        if n_times_after_sinc_filter < pool_size:
+            raise ValueError(
+                f"n_times ({self.n_times}) leaves only {n_times_after_sinc_filter} "
+                f"samples after the temporal sinc filter, fewer than pool_size "
+                f"({pool_size}): the average pooling would produce no temporal "
+                "position. Increase n_times, or decrease pool_size (and "
+                "time_filter_len for padding='valid')."
+            )
+
         size_after_pooling = (
             (n_times_after_sinc_filter - pool_size) // pool_stride
         ) + 1

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/tmsanet.py

@@ -141,7 +141,18 @@ class TMSANet(EEGModuleMixin, nn.Module, license="mit"):
         )
         # The reference flattens (embed_dim, tokens), embedding-major.
         self.flatten = Rearrange("b n d -> b (d n)")
+        if self.n_times is None:
+            raise ValueError(
+                "TMSANet needs to know the input length: pass n_times "
+                "(or input_window_seconds and sfreq)."
+            )
         n_tokens = (self.n_times - pool_size) // pool_stride + 1
+        if n_tokens < 1:
+            raise ValueError(
+                f"n_times ({self.n_times}) must be at least pool_size ({pool_size}) "
+                f"so the average pooling keeps at least one temporal token "
+                f"(would keep {n_tokens})."
+            )
         self.final_layer = nn.Linear(embed_dim * n_tokens, self.n_outputs)
 
     def forward(self, x: Tensor) -> Tensor:

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/models/zuna.py

@@ -138,9 +138,11 @@ class ZUNA(EEGModuleMixin, nn.Module, license="apache-2.0"):
     latent_dim : int, optional
         Per-token output dimension of the encoder. The default is ``32``.
     max_seqlen : int, optional
-        Length of the rotary frequency table. It must be at least
-        ``max(pos_bins, the number of temporal patches)``. The default is
-        ``256``.
+        Minimum length of the rotary frequency table. The table is extended to
+        ``max(pos_bins, the number of temporal patches)`` when the window needs
+        more rows, so longer windows than the release's 8192 samples are
+        accepted (positions past the pretraining range are extrapolated by the
+        rotary embedding). The default is ``256``.
     rope_theta : float, optional
         Base period of the rotary positional embedding. The default is
         ``10000.0``.
@@ -263,12 +265,6 @@ class ZUNA(EEGModuleMixin, nn.Module, license="apache-2.0"):
             raise ValueError("ZUNA requires finite 3D locations for every channel.")
         if head_dim % 8 != 0:
             raise ValueError("head_dim must be divisible by eight for 4D RoPE.")
-        if coarse_time_points > max_seqlen:
-            raise ValueError(
-                f"n_times ({self.n_times}) gives {coarse_time_points} patches of "
-                f"{fine_time_pts} samples; max_seqlen={max_seqlen} allows at most "
-                f"{max_seqlen * fine_time_pts} samples."
-            )
 
         # Layers
         self.patch_embedding = nn.Sequential(
@@ -298,7 +294,11 @@ class ZUNA(EEGModuleMixin, nn.Module, license="apache-2.0"):
         )
 
         rotary_frequency_table = _build_rotary_frequency_table(
-            torch.arange(max_seqlen, dtype=torch.float32),
+            # RoPE is analytic and the tables are non-persistent buffers: extend the
+            # table to the window instead of capping it (rows < max_seqlen unchanged).
+            torch.arange(
+                max(max_seqlen, pos_bins, coarse_time_points), dtype=torch.float32
+            ),
             axis_dim=rotary_axis_dim,
             theta=rope_theta,
         )

{braindecode-1.8.1.dev180761562 → braindecode-1.8.1.dev180835552}/braindecode/modules/attention.py

@@ -1004,10 +1004,13 @@ class MultiHeadAttention(nn.Module):
         x : Tensor
             Input tensor of shape ``(batch, seq, emb_size)``.
         mask : Tensor, optional
-            Attention mask following PyTorch SDPA convention: for boolean
-            masks ``True`` means **ignore** that position; for float
-            masks the values are **added** to attention scores before
-            softmax.
+            Attention mask following the convention of
+            :func:`torch.nn.functional.scaled_dot_product_attention`: for
+            boolean masks ``True`` means the position **takes part** in
+            attention and ``False`` that it is ignored (the opposite of the
+            ``key_padding_mask`` of :class:`torch.nn.MultiheadAttention`);
+            for float masks the values are **added** to the attention scores
+            before softmax.
         """
         queries = self.rearrange_stack(self.queries(x))
         keys = self.rearrange_stack(self.keys(x))
@@ -1086,14 +1089,16 @@ class CrissCrossTransformerEncoderLayer(nn.Module):
             nhead // 2,
             dropout=dropout,
             bias=bias,
-            batch_first=batch_first,
+            # _sa_block always builds (batch, sequence, features) inputs
+            batch_first=True,
         )
         self.self_attn_t = nn.MultiheadAttention(
             d_model // 2,
             nhead // 2,
             dropout=dropout,
             bias=bias,
-            batch_first=batch_first,
+            # _sa_block always builds (batch, sequence, features) inputs
+            batch_first=True,
         )
 
         # Implementation of Feedforward model
@@ -1101,6 +1106,9 @@ class CrissCrossTransformerEncoderLayer(nn.Module):
         self.dropout = nn.Dropout(dropout)
         self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias)
 
+        # kept for the signature of nn.TransformerEncoderLayer: the input is
+        # always (batch, channels, n_patches, patch_size)
+        self.batch_first = batch_first
         self.norm_first = norm_first
         self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
         self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
@@ -1116,10 +1124,17 @@ class CrissCrossTransformerEncoderLayer(nn.Module):
         is_causal: bool = False,
     ) -> Tensor:
         x = src
-        x = x + self._sa_block(
-            self.norm1(x), src_mask, src_key_padding_mask, is_causal=is_causal
-        )
-        x = x + self._ff_block(self.norm2(x))
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x), src_mask, src_key_padding_mask, is_causal=is_causal
+            )
+            x = x + self._ff_block(self.norm2(x))
+        else:
+            x = self.norm1(
+                x
+                + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal)
+            )
+            x = self.norm2(x + self._ff_block(x))
         return x
 
     # self-attention block