braindecode 1.8.1.dev173387713__tar.gz → 1.8.1.dev174857482__tar.gz

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/NOTICE.txt

@@ -43,6 +43,8 @@ The following components are licensed under the MIT License:
 - `braindecode/models/tcformer.py`
 - `braindecode/models/ifnet.py`
 - `braindecode/models/mirepnet.py`
+- `braindecode/models/popt.py`
+- `braindecode/modules/popt_modules.py`
 
 ### Apache-2.0 Licensed Files
 

{braindecode-1.8.1.dev173387713/braindecode.egg-info → braindecode-1.8.1.dev174857482}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev173387713
+Version: 1.8.1.dev174857482
 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>
@@ -60,6 +60,8 @@ Provides-Extra: typing
 Requires-Dist: exca>=0.5.10; extra == "typing"
 Requires-Dist: numpydantic>=1.7; extra == "typing"
 Provides-Extra: docs
+Requires-Dist: pytest<9.1; extra == "docs"
+Requires-Dist: sphinx-markdown-builder==0.6.11; extra == "docs"
 Requires-Dist: sphinx_gallery>=0.17; extra == "docs"
 Requires-Dist: sphinx_rtd_theme; extra == "docs"
 Requires-Dist: sphinx-autodoc-typehints; extra == "docs"

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/models/__init__.py

@@ -48,6 +48,7 @@ from .msvtnet import MSVTNet
 from .mvpformer import MVPFormer
 from .neuropose import NeuroPose
 from .patchedtransformer import PBT
+from .popt import PopulationTransformer
 from .reve import REVE
 from .sccnet import SCCNet
 from .sensingdynamics import SensingDynamics
@@ -144,6 +145,7 @@ __all__ = [
     "MSVTNet",
     "MVPFormer",
     "PBT",
+    "PopulationTransformer",
     "REVE",
     "SCCNet",
     "ShallowFBCSPNet",

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/models/deep4.py

@@ -186,13 +186,18 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         # When padronize all layers,
         # add the old's parameters here
         self.mapping = {
-            "conv_time.weight": "conv_time_spat.conv_time.weight",
-            "conv_spat.weight": "conv_time_spat.conv_spat.weight",
-            "conv_time.bias": "conv_time_spat.conv_time.bias",
-            "conv_spat.bias": "conv_time_spat.conv_spat.bias",
             "conv_classifier.weight": "final_layer.conv_classifier.weight",
             "conv_classifier.bias": "final_layer.conv_classifier.bias",
         }
+        if self.split_first_layer:
+            self.mapping.update(
+                {
+                    "conv_time.weight": "conv_time_spat.conv_time.weight",
+                    "conv_spat.weight": "conv_time_spat.conv_spat.weight",
+                    "conv_time.bias": "conv_time_spat.conv_time.bias",
+                    "conv_spat.bias": "conv_time_spat.conv_spat.bias",
+                }
+            )
 
         if self.stride_before_pool:
             conv_stride = self.pool_time_stride
@@ -215,6 +220,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
                     filter_time_length=filter_time_length,
                     bias_time=True,
                     bias_spat=not self.batch_norm,
+                    stride=(conv_stride, 1),
                 ),
             )
             n_filters_conv = self.n_filters_spat
@@ -318,14 +324,16 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
 
         # Initialization, xavier is same as in our paper...
         # was default from lasagne
-        init.xavier_uniform_(self.conv_time_spat.conv_time.weight, gain=1)
-        # maybe no bias in case of no split layer and batch norm
-        if self.split_first_layer or (not self.batch_norm):
-            init.constant_(self.conv_time_spat.conv_time.bias, 0)
         if self.split_first_layer:
+            init.xavier_uniform_(self.conv_time_spat.conv_time.weight, gain=1)
+            init.constant_(self.conv_time_spat.conv_time.bias, 0)
             init.xavier_uniform_(self.conv_time_spat.conv_spat.weight, gain=1)
             if not self.batch_norm:
                 init.constant_(self.conv_time_spat.conv_spat.bias, 0)
+        else:
+            init.xavier_uniform_(self.conv_time.weight, gain=1)
+            if not self.batch_norm:
+                init.constant_(self.conv_time.bias, 0)
         if self.batch_norm:
             init.constant_(self.bnorm.weight, 1)
             init.constant_(self.bnorm.bias, 0)

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/models/eegpt.py

@@ -88,7 +88,8 @@ class EEGPT(EEGModuleMixin, nn.Module):
         - *Role.* Maps raw spatio-temporal EEG patches into a sequence of latent tokens :math:`z`.
     - `EEGPT.chans_id` **(Channel Identification)**
         - *Operations.* A buffer containing channel indices mapped from the standard channel names provided
-          in ``chs_info`` [eegpt]_.
+          in ``chs_info`` [eegpt]_. It is rebuilt from the montage at construction and not stored in
+          checkpoints, so pretrained weights load on any montage.
         - *Role.* Provides the spatial identity for each input channel, allowing the model to look up
           the correct channel embedding vector :math:`\varsigma_i`.
     - **Local Spatio-Temporal Embedding** (Input Processing)
@@ -338,8 +339,12 @@ class EEGPT(EEGModuleMixin, nn.Module):
         else:
             self.channel_names = None  # type: ignore
 
+        # Derived from the montage above, so it is not saved in the state dict:
+        # loading a checkpoint must not replace the IDs of this model's channels.
         self.register_buffer(
-            "chans_id", self.target_encoder.prepare_chan_ids(self.channel_names)
+            "chans_id",
+            self.target_encoder.prepare_chan_ids(self.channel_names),
+            persistent=False,
         )
 
         self.flattened_encoder_output_dim = (
@@ -391,6 +396,13 @@ class EEGPT(EEGModuleMixin, nn.Module):
             "embed_dim": self.embed_dim,
         }
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # Checkpoints saved before ``chans_id`` became non-persistent (including
+        # the released EEGPT weights) still carry it. Drop it, so the IDs keep
+        # matching this model's montage and a different montage can load.
+        state_dict.pop(prefix + "chans_id", None)
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     def forward(self, x, return_features=False):
         """
         Forward pass.

braindecode-1.8.1.dev174857482/braindecode/models/popt.py

@@ -0,0 +1,320 @@
+# Authors: Christopher Wang, Geeling Chau (original implementation)
+#          Adam Mounir <am91ris@gmail.com> (braindecode adaptation)
+#
+# License: MIT
+# Adapted from https://github.com/czlwang/PopulationTransformer
+"""PopulationTransformer (PopT, Chau et al. 2024).
+
+* paper: https://arxiv.org/abs/2406.03044
+* code: https://github.com/czlwang/PopulationTransformer
+* weights: https://huggingface.co/PopulationTransformer/popt_brainbert_stft
+"""
+
+from __future__ import annotations
+
+import warnings
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.models.util import extract_channel_locations_from_chs_info
+from braindecode.modules.popt_modules import (
+    _PopTInputEmbedding,
+    _PopTSpecPredictionHead,
+)
+
+
+class PopulationTransformer(EEGModuleMixin, nn.Module, license="mit"):
+    r"""PopulationTransformer (PopT) from Chau et al. (2024) [PopT2024]_.
+
+    :bdg-danger:`Foundation Model` :bdg-info:`Attention/Transformer`
+
+    PopT is a self-supervised **population** model for intracranial recordings
+    (sEEG/iEEG). It does not encode a raw time signal; instead each electrode is
+    represented by a feature vector — typically the frozen embedding of a
+    per-channel foundation model such as :class:`~braindecode.models.BrainBERT` —
+    and PopT aggregates across electrodes. Every electrode feature is linearly
+    projected and given a fixed sinusoidal **spatial** position encoding built
+    from its integer anatomical coordinates (one embedding per X/Y/Z axis plus a
+    sequence id). A ``CLS`` token is prepended, a stack of standard Transformer
+    encoder layers mixes the population, and the ``CLS`` output is the pooled
+    representation used for downstream decoding. Pre-training is by masked /
+    replaced-token modelling over the electrode population.
+
+    Following the braindecode convention, the per-electrode feature vector plays
+    the role of the ``n_times`` axis, so the model keeps the standard
+    ``(batch, n_chans, n_times)`` input signature: ``n_chans`` is the number of
+    electrodes and ``n_times`` is the upstream feature dimension (768 for
+    BrainBERT ``stft`` features). Electrode coordinates are read from
+    ``chs_info`` (their ``loc``) and discretised to absolute integer indices
+    inside the model, as upstream feeds them; when no positions are available
+    the electrodes fall back to distinct sequential indices.
+
+    The ``CLS`` output goes through a single linear layer, as in the upstream
+    fine-tuning model (``PtDownstreamModel.linear_out``, one logit trained with
+    binary cross-entropy there; ``n_outputs=1`` reproduces it).
+
+    The defaults are the released ``popt_brainbert_stft`` configuration:
+    ``hidden_dim=512``, ``ffn_dim=2048``, ``n_heads=8``, ``n_layers=6``, used
+    on ``n_times=768`` BrainBERT features (~20M parameters).
+
+    .. important::
+       **Pre-trained weights available.** The official checkpoint is released by
+       the authors and loads directly::
+
+           model = PopulationTransformer.from_pretrained(
+               "braindecode/popt-pretrained", n_outputs=2
+           )
+
+       It uses the default configuration; ``n_chans`` and ``n_outputs``
+       may be changed freely, as the population is pooled through the ``CLS``
+       token and the classification head is task-specific (the checkpoint
+       carries no trained fine-tuning head).
+
+    .. warning::
+       **Evaluate with a time-blocked split.** The paper's downstream results
+       use a random 80/10/10 split over word-aligned 5 s windows. Words are a
+       fraction of a second apart, so almost every test window overlaps a
+       training window, and labels that drift slowly in time leak into
+       training. Re-running the paper setup (7 subjects, 3 seeds) with
+       contiguous blocks of time, and dropping training windows that overlap
+       the test set, pretrained PopT goes from 0.79 to 0.51 ROC-AUC on Pitch
+       (chance) and from 0.89 to 0.64 on Volume. Onset (0.86 to 0.84) and
+       Speech (0.90 to 0.84) hold up, and pretraining still beats training
+       from scratch on Onset, Speech and Volume. The model and weights are
+       not affected; the issue is only in the evaluation. When fine-tuning,
+       split by blocks of time.
+
+    .. versionadded:: 1.8.2
+
+    Parameters
+    ----------
+    hidden_dim : int, optional
+        Transformer model width ``D``. Must be divisible by 8. Default 512, as
+        the released model.
+    ffn_dim : int, optional
+        Inner dimension of the Transformer feed-forward blocks. Default 2048, as
+        the released model.
+    n_layers : int, optional
+        Number of Transformer encoder layers. Default 6, as the released model.
+    n_heads : int, optional
+        Number of attention heads. Default 8, as the released model.
+    max_len : int, optional
+        Size of the coordinate table (largest addressable integer coordinate).
+        Default 5000, as upstream.
+    coord_units : {"m", "raw"}, optional
+        How ``chs_info`` positions become integer coordinates. ``"m"`` (default)
+        treats them as MNE metres and rounds them to millimetres. ``"raw"``
+        rounds the positions as they are: use it when ``x/y/z`` already hold the
+        Brain Treebank integer (left, inferior, posterior) coordinates, as NEMAR
+        nm000253 stores them. Either way the indices are absolute, not
+        shifted, as upstream feeds them (``pt_supervised_task_coords.py``).
+        They match the pretrained checkpoint only if the positions are already
+        in the upstream (left, inferior, posterior) space; MNE head-frame
+        positions (e.g. a standard montage) are **not** that space. Indices
+        outside ``[0, max_len - 1]`` are clamped, with a warning. You can also
+        pass ``coords`` to :meth:`forward` directly.
+    shift_coords : bool, optional
+        If ``True``, shift each axis so that its smallest index is 0. Default
+        ``False``. Upstream does not shift, and the shift changes the position
+        encoding, and so the output of the pretrained model; use it only for
+        positions with negative values when training from scratch.
+    activation : type[nn.Module], optional
+        Feed-forward activation, given as a class. Default :class:`~torch.nn.GELU`.
+    drop_prob : float, optional
+        Dropout probability. Default 0.1.
+
+    References
+    ----------
+    .. [PopT2024] Chau, G., Wang, C., Talukder, S., Subramaniam, V., Soedarmadji,
+       S., Yue, Y., Katz, B., & Barbu, A. (2024). Population Transformer:
+       Learning Population-level Representations of Neural Activity. arXiv
+       preprint arXiv:2406.03044.
+    """
+
+    def __init__(
+        self,
+        # braindecode parameters
+        n_outputs=None,
+        n_chans=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        # model-specific parameters
+        *,
+        hidden_dim: int = 512,
+        ffn_dim: int = 2048,
+        n_layers: int = 6,
+        n_heads: int = 8,
+        max_len: int = 5000,
+        coord_units: str = "m",
+        shift_coords: bool = False,
+        activation: type[nn.Module] = nn.GELU,
+        drop_prob: float = 0.1,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+        if coord_units not in ("m", "raw"):
+            raise ValueError(f"coord_units must be 'm' or 'raw', got {coord_units!r}.")
+
+        self.input_dim = self.n_times
+        self.hidden_dim = hidden_dim
+        self.ffn_dim = ffn_dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.max_len = max_len
+        self.coord_units = coord_units
+        self.shift_coords = shift_coords
+
+        self.input_embedding = _PopTInputEmbedding(
+            input_dim=self.input_dim,
+            hidden_dim=hidden_dim,
+            max_len=max_len,
+            drop_prob=drop_prob,
+        )
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim,
+            nhead=n_heads,
+            dim_feedforward=ffn_dim,
+            activation=activation(),
+            dropout=drop_prob,
+            batch_first=True,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=n_layers
+        )
+        # Pre-training head, unused here; kept so the checkpoint loads strictly.
+        self.spec_prediction_head = _PopTSpecPredictionHead(hidden_dim, self.input_dim)
+        self.final_layer = nn.Linear(hidden_dim, self.n_outputs)
+
+        self.register_buffer(
+            "electrode_coords", self._coords_from_chs_info(), persistent=False
+        )
+
+    def _coords_from_chs_info(self) -> torch.Tensor:
+        """Integer ``(n_chans, 3)`` coordinates from ``chs_info``, or sequential ones."""
+        chs_info = getattr(self, "_chs_info", None)
+        loc = extract_channel_locations_from_chs_info(
+            chs_info, num_channels=self.n_chans
+        )
+        # The helper does not screen NaNs or partial montages.
+        if loc is None or loc.shape[0] != self.n_chans or not np.isfinite(loc).all():
+            idx = torch.arange(self.n_chans, dtype=torch.long)
+            coords = idx.unsqueeze(1).repeat(1, 3)
+        else:
+            loc_t = torch.as_tensor(loc, dtype=torch.float)
+            if self.coord_units == "m":
+                loc_t = loc_t * 1000.0
+            coords = loc_t.round().long()
+            if self.shift_coords:
+                coords = coords - coords.min(dim=0, keepdim=True).values
+            if bool(((coords < 0) | (coords >= self.max_len)).any()):
+                hint = (
+                    "raise `max_len`"
+                    if self.shift_coords
+                    else "set `shift_coords=True` to train from scratch"
+                )
+                warnings.warn(
+                    "Some electrode coordinates from chs_info fall outside "
+                    f"[0, {self.max_len - 1}] and are clamped. The pretrained "
+                    "model expects the non-negative integer (left, inferior, "
+                    "posterior) indices of the upstream data; pass `coords` to "
+                    f"forward, or {hint}.",
+                    UserWarning,
+                    stacklevel=4,
+                )
+        return coords.clamp(0, self.max_len - 1)
+
+    def load_state_dict(self, state_dict, *args, **kwargs):
+        """Also accept the untrained ``final_layer.{norm,fc}`` head of the HF mirror."""
+        remapped = OrderedDict()
+        for key, value in state_dict.items():
+            if key.startswith("final_layer.norm."):
+                continue
+            if key.startswith("final_layer.fc."):
+                key = "final_layer." + key[len("final_layer.fc.") :]
+            remapped[key] = value
+        return super().load_state_dict(remapped, *args, **kwargs)
+
+    def reset_head(self, n_outputs: int) -> None:
+        """Swap the classification head for a new number of outputs."""
+        old = self.final_layer.weight
+        self._set_n_outputs(n_outputs)
+        self.final_layer = nn.Linear(self.hidden_dim, n_outputs).to(
+            device=old.device, dtype=old.dtype
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        coords: torch.Tensor | None = None,
+        seq_id: torch.Tensor | None = None,
+        return_features: bool = False,
+        key_padding_mask: torch.Tensor | None = None,
+    ):
+        """Aggregate a population of electrode features.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Per-electrode features of shape ``(batch, n_chans, n_times)``, where
+            ``n_times`` is the upstream feature dimension.
+        coords : torch.Tensor, optional
+            Integer coordinates of shape ``(batch, n_chans, 3)``. Defaults to the
+            coordinates derived from ``chs_info`` at construction, broadcast over
+            the batch.
+        seq_id : torch.Tensor, optional
+            Integer sequence ids of shape ``(batch, n_chans)``. Defaults to zero
+            (single population).
+        return_features : bool
+            If ``True``, return ``{"features": cls, "cls_token": cls}`` (the
+            pooled ``CLS`` representation) instead of the class logits.
+        key_padding_mask : torch.Tensor, optional
+            Boolean ``(batch, n_chans)`` mask, ``True`` for padded electrodes, so
+            recordings with different electrode sets can share a batch (upstream
+            ``src_key_padding_mask``). The ``CLS`` token is never masked.
+
+        Returns
+        -------
+        torch.Tensor or dict
+            Class logits of shape ``(batch, n_outputs)``, or the feature dict
+            when ``return_features`` is set.
+        """
+        batch_size, n_chans, _ = x.shape
+        if coords is None:
+            n_known = self.electrode_coords.shape[0]
+            if n_chans != n_known:
+                raise ValueError(
+                    f"x has {n_chans} electrodes but the model has coordinates "
+                    f"for {n_known}; pass `coords` of shape (batch, {n_chans}, 3)."
+                )
+            coords = self.electrode_coords.unsqueeze(0).expand(batch_size, -1, -1)
+        if seq_id is None:
+            seq_id = torch.zeros(batch_size, n_chans, dtype=torch.long, device=x.device)
+
+        h = self.input_embedding(x, coords, seq_id)
+        # New name (not reassigning the argument) keeps TorchScript's narrowing.
+        padded_mask: torch.Tensor | None = None
+        if key_padding_mask is not None:
+            cls_keep = torch.zeros(batch_size, 1, dtype=torch.bool, device=x.device)
+            padded_mask = torch.cat([cls_keep, key_padding_mask.to(torch.bool)], dim=1)
+        z = self.transformer_encoder(h, src_key_padding_mask=padded_mask)
+        cls_token = z[:, 0, :]
+        logits = self.final_layer(cls_token)
+        if return_features:
+            # Scripted forward stays monomorphic (same pattern as Brant).
+            if torch.jit.is_scripting():
+                return logits
+            return {"features": cls_token, "cls_token": cls_token}
+        return logits

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/models/shallow_fbcsp.py

@@ -124,13 +124,18 @@ class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
         self.drop_prob = drop_prob
 
         self.mapping = {
-            "conv_time.weight": "conv_time_spat.conv_time.weight",
-            "conv_spat.weight": "conv_time_spat.conv_spat.weight",
-            "conv_time.bias": "conv_time_spat.conv_time.bias",
-            "conv_spat.bias": "conv_time_spat.conv_spat.bias",
             "conv_classifier.weight": "final_layer.conv_classifier.weight",
             "conv_classifier.bias": "final_layer.conv_classifier.bias",
         }
+        if self.split_first_layer:
+            self.mapping.update(
+                {
+                    "conv_time.weight": "conv_time_spat.conv_time.weight",
+                    "conv_spat.weight": "conv_time_spat.conv_spat.weight",
+                    "conv_time.bias": "conv_time_spat.conv_time.bias",
+                    "conv_spat.bias": "conv_time_spat.conv_spat.bias",
+                }
+            )
 
         self.add_module("ensuredims", Ensure4d())
         pool_class = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)[self.pool_mode]
@@ -203,14 +208,16 @@ class ShallowFBCSPNet(EEGModuleMixin, nn.Sequential):
         self.add_module("final_layer", module)
 
         # Initialization, xavier is same as in paper...
-        init.xavier_uniform_(self.conv_time_spat.conv_time.weight, gain=1)
-        # maybe no bias in case of no split layer and batch norm
-        if self.split_first_layer or (not self.batch_norm):
-            init.constant_(self.conv_time_spat.conv_time.bias, 0)
         if self.split_first_layer:
+            init.xavier_uniform_(self.conv_time_spat.conv_time.weight, gain=1)
+            init.constant_(self.conv_time_spat.conv_time.bias, 0)
             init.xavier_uniform_(self.conv_time_spat.conv_spat.weight, gain=1)
             if not self.batch_norm:
                 init.constant_(self.conv_time_spat.conv_spat.bias, 0)
+        else:
+            init.xavier_uniform_(self.conv_time.weight, gain=1)
+            if not self.batch_norm:
+                init.constant_(self.conv_time.bias, 0)
         if self.batch_norm:
             init.constant_(self.bnorm.weight, 1)
             init.constant_(self.bnorm.bias, 0)

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/models/summary.csv

@@ -52,6 +52,7 @@ IFNet,Motor Imagery,Prediction,250,"n_chans, n_outputs, n_times, sfreq",9860,"IF
 BrainModule,Speech Decoding,Prediction,250,"n_chans, n_outputs, n_times, sfreq",6186909,"BrainModule(n_chans=64, n_outputs=29, n_times=160, sfreq=1000)","Convolution","EEG, MEG"
 Brant,Epilepsy,Prediction,250,"n_chans, n_outputs, n_times, sfreq",2902780,"Brant(n_chans=16, n_outputs=2, n_times=3000, sfreq=250)","Foundation Model,Attention/Transformer",iEEG
 PBT,General,Prediction,250,"n_chans, n_outputs, n_times",818948,"PBT(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Foundation Model",EEG
+PopulationTransformer,General,"Prediction, Embedding",2048,"n_chans, n_outputs, n_times",19967746,"PopulationTransformer(n_chans=8, n_outputs=2, n_times=768)","Foundation Model,Attention/Transformer",iEEG
 SSTDPN,Motor Imagery,Prediction,250,"n_chans, n_outputs, n_times",19502,"SSTDPN(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer",EEG
 BENDR,General,"Prediction, Embedding",250,"n_chans, n_times, n_outputs",157141049,"BENDR(n_chans=22, n_outputs=4, n_times=1000)","Foundation Model,Convolution",EEG
 LUNA,General,"Prediction, Embedding",128,"n_chans, n_times, sfreq, chs_info",7100731,"LUNA(n_chans=22, n_times=512, sfreq=128)","Convolution,Channel,Foundation Model",EEG

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/models/util.py

@@ -421,6 +421,7 @@ models_mandatory_parameters: list[
         ["n_chans", "n_outputs", "n_times"],
         {"n_times": 2048, "sfreq": 2048.0},  # the pretraining rate, 1 s windows
     ),
+    ("PopulationTransformer", ["n_chans", "n_outputs", "n_times"], None),
     (
         "InterpolatedBIOT",
         ["chs_info", "n_outputs", "sfreq", "n_times"],

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/modules/convolution.py

@@ -48,6 +48,14 @@ class AvgPool2dWithConv(nn.Module):
         # that initializes parameters or something
         self._pool_weights = None
 
+    def _pool_weights_match_input(self, x, weight_shape):
+        return (
+            self._pool_weights is not None
+            and tuple(self._pool_weights.size()) == tuple(weight_shape)
+            and self._pool_weights.device == x.device
+            and self._pool_weights.dtype == x.dtype
+        )
+
     def forward(self, x):
         # Create weights for the convolution on demand:
         # size or type of x changed...
@@ -58,17 +66,10 @@ class AvgPool2dWithConv(nn.Module):
             self.kernel_size[0],
             self.kernel_size[1],
         )
-        if self._pool_weights is None or (
-            (tuple(self._pool_weights.size()) != tuple(weight_shape))
-            or (self._pool_weights.is_cuda != x.is_cuda)
-            or (self._pool_weights.data.type() != x.data.type())
-        ):
+        if not self._pool_weights_match_input(x, weight_shape):
             n_pool = np.prod(self.kernel_size)
             weights = np_to_th(np.ones(weight_shape, dtype=np.float32) / float(n_pool))
-            weights = weights.type_as(x)
-            if x.is_cuda:
-                weights = weights.cuda()
-            self._pool_weights = weights
+            self._pool_weights = weights.to(device=x.device, dtype=x.dtype)
 
         pooled = F.conv2d(
             x,
@@ -143,6 +144,9 @@ class CombinedConv(nn.Module):
         Whether to use bias in the temporal conv
     bias_spat: bool
         Whether to use bias in the spatial conv
+    stride: int | tuple[int, int]
+        Stride of the merged convolution. This is applied at the spatial
+        convolution stage, matching the unmerged temporal-then-spatial form.
 
     Examples
     --------
@@ -164,6 +168,7 @@ class CombinedConv(nn.Module):
         filter_time_length=25,
         bias_time=True,
         bias_spat=True,
+        stride=1,
     ):
         super().__init__()
         self.bias_time = bias_time
@@ -172,7 +177,11 @@ class CombinedConv(nn.Module):
             1, n_filters_time, (filter_time_length, 1), bias=bias_time, stride=1
         )
         self.conv_spat = nn.Conv2d(
-            n_filters_time, n_filters_spat, (1, in_chans), bias=bias_spat, stride=1
+            n_filters_time,
+            n_filters_spat,
+            (1, in_chans),
+            bias=bias_spat,
+            stride=stride,
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -210,7 +219,9 @@ class CombinedConv(nn.Module):
 
         bias = calculated_bias
 
-        return F.conv2d(x, weight=combined_weight, bias=bias, stride=(1, 1))
+        return F.conv2d(
+            x, weight=combined_weight, bias=bias, stride=self.conv_spat.stride
+        )
 
 
 class CausalConv1d(nn.Conv1d):
@@ -274,16 +285,16 @@ class CausalConv1d(nn.Conv1d):
         )
 
     def forward(self, X):
-        out = F.conv1d(
+        X = F.pad(X, (self.padding[0], 0))
+        return F.conv1d(
             X,
             self.weight,
             self.bias,
             stride=self.stride,
-            padding=self.padding,
+            padding=0,
             dilation=self.dilation,
             groups=self.groups,
         )
-        return out[..., : -self.padding[0]]
 
 
 class DepthwiseConv2d(torch.nn.Conv2d):

{braindecode-1.8.1.dev173387713 → braindecode-1.8.1.dev174857482}/braindecode/modules/layers.py

@@ -62,6 +62,8 @@ class Chomp1d(nn.Module):
         return "chomp_size={}".format(self.chomp_size)
 
     def forward(self, x):
+        if self.chomp_size == 0:
+            return x.contiguous()
         return x[:, :, : -self.chomp_size].contiguous()
 
 
@@ -110,8 +112,8 @@ class TimeDistributed(nn.Module):
             Shape (batch_size, seq_len, output_size).
         """
         b, s, c, t = x.shape
-        out = self.module(x.view(b * s, c, t))
-        return out.view(b, s, -1)
+        out = self.module(x.reshape(b * s, c, t))
+        return out.reshape(b, s, -1)
 
 
 class DropPath(nn.Module):