braindecode 1.2.0.dev175337561__tar.gz → 1.2.0.dev176281713__tar.gz

{braindecode-1.2.0.dev175337561/braindecode.egg-info → braindecode-1.2.0.dev176281713}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.2.0.dev175337561
+Version: 1.2.0.dev176281713
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -49,6 +49,9 @@ Requires-Dist: mypy; extra == "tests"
 Provides-Extra: docs
 Requires-Dist: sphinx_gallery; extra == "docs"
 Requires-Dist: sphinx_rtd_theme; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
+Requires-Dist: sphinx-autobuild; extra == "docs"
+Requires-Dist: sphinxcontrib-bibtex; extra == "docs"
 Requires-Dist: pydata_sphinx_theme; extra == "docs"
 Requires-Dist: numpydoc; extra == "docs"
 Requires-Dist: memory_profiler; extra == "docs"
@@ -59,6 +62,7 @@ Requires-Dist: lightning; extra == "docs"
 Requires-Dist: seaborn; extra == "docs"
 Requires-Dist: pre-commit; extra == "docs"
 Requires-Dist: openneuro-py; extra == "docs"
+Requires-Dist: plotly; extra == "docs"
 Provides-Extra: all
 Requires-Dist: braindecode[docs,moabb,tests]; extra == "all"
 Dynamic: license-file

braindecode-1.2.0.dev176281713/braindecode/datasets/experimental.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import Callable, Sequence
+
+import mne_bids
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+class BIDSIterableDataset(IterableDataset):
+    """Dataset for loading BIDS.
+
+    .. warning::
+        This class is experimental and may change in the future.
+
+    .. warning::
+        This dataset is not consistent with the Braindecode API.
+
+    This class has the same parameters as the :func:`mne_bids.find_matching_paths` function
+    as it will be used to find the files to load. The default ``extensions`` parameter was changed.
+
+    More information on BIDS (Brain Imaging Data Structure)
+    can be found at https://bids.neuroimaging.io
+
+    Examples
+    --------
+    >>> from braindecode.datasets import BaseDataset, BaseConcatDataset
+    >>> from braindecode.datasets.bids import BIDSIterableDataset, _description_from_bids_path
+    >>> from braindecode.preprocessing import create_fixed_length_windows
+    >>>
+    >>> def my_reader_fn(path):
+    ...     raw = mne_bids.read_raw_bids(path)
+    ...     desc = _description_from_bids_path(path)
+    ...     ds = BaseDataset(raw, description=desc)
+    ...     windows_ds = create_fixed_length_windows(
+    ...         BaseConcatDataset([ds]),
+    ...         window_size_samples=400,
+    ...         window_stride_samples=200,
+    ...     )
+    ...     return windows_ds
+    >>>
+    >>> dataset = BIDSIterableDataset(
+    ...     reader_fn=my_reader_fn,
+    ...     root="root/of/my/bids/dataset/",
+    ... )
+
+    Parameters
+    ----------
+    reader_fn : Callable[[mne_bids.BIDSPath], Sequence]
+        A function that takes a BIDSPath and returns a dataset.
+    pool_size : int
+        The number of recordings to read and sample from.
+    bids_paths : list[mne_bids.BIDSPath] | None
+        A list of BIDSPaths to load. If None, will use the paths found by
+        :func:`mne_bids.find_matching_paths` and the arguments below.
+    root : pathlib.Path | str
+        The root of the BIDS path.
+    subjects : str | array-like of str | None
+        The subject ID. Corresponds to "sub".
+    sessions : str | array-like of str | None
+        The acquisition session. Corresponds to "ses".
+    tasks : str | array-like of str | None
+        The experimental task. Corresponds to "task".
+    acquisitions: str | array-like of str | None
+        The acquisition parameters. Corresponds to "acq".
+    runs : str | array-like of str | None
+        The run number. Corresponds to "run".
+    processings : str | array-like of str | None
+        The processing label. Corresponds to "proc".
+    recordings : str | array-like of str | None
+        The recording name. Corresponds to "rec".
+    spaces : str | array-like of str | None
+        The coordinate space for anatomical and sensor location
+        files (e.g., ``*_electrodes.tsv``, ``*_markers.mrk``).
+        Corresponds to "space".
+        Note that valid values for ``space`` must come from a list
+        of BIDS keywords as described in the BIDS specification.
+    splits : str | array-like of str | None
+        The split of the continuous recording file for ``.fif`` data.
+        Corresponds to "split".
+    descriptions : str | array-like of str | None
+        This corresponds to the BIDS entity ``desc``. It is used to provide
+        additional information for derivative data, e.g., preprocessed data
+        may be assigned ``description='cleaned'``.
+    suffixes : str | array-like of str | None
+        The filename suffix. This is the entity after the
+        last ``_`` before the extension. E.g., ``'channels'``.
+        The following filename suffix's are accepted:
+        'meg', 'markers', 'eeg', 'ieeg', 'T1w',
+        'participants', 'scans', 'electrodes', 'coordsystem',
+        'channels', 'events', 'headshape', 'digitizer',
+        'beh', 'physio', 'stim'
+    extensions : str | array-like of str | None
+        The extension of the filename. E.g., ``'.json'``.
+        By default, uses the ones accepted by :func:`mne_bids.read_raw_bids`.
+    datatypes : str | array-like of str | None
+        The BIDS data type, e.g., ``'anat'``, ``'func'``, ``'eeg'``, ``'meg'``,
+        ``'ieeg'``.
+    check : bool
+        If ``True``, only returns paths that conform to BIDS. If ``False``
+        (default), the ``.check`` attribute of the returned
+        :class:`mne_bids.BIDSPath` object will be set to ``True`` for paths that
+        do conform to BIDS, and to ``False`` for those that don't.
+    preload : bool
+        If True, preload the data. Defaults to False.
+    n_jobs : int
+        Number of jobs to run in parallel. Defaults to 1.
+    """
+
+    def __init__(
+        self,
+        reader_fn: Callable[[mne_bids.BIDSPath], Sequence],
+        pool_size: int = 4,
+        bids_paths: list[mne_bids.BIDSPath] | None = None,
+        root: Path | str | None = None,
+        subjects: str | list[str] | None = None,
+        sessions: str | list[str] | None = None,
+        tasks: str | list[str] | None = None,
+        acquisitions: str | list[str] | None = None,
+        runs: str | list[str] | None = None,
+        processings: str | list[str] | None = None,
+        recordings: str | list[str] | None = None,
+        spaces: str | list[str] | None = None,
+        splits: str | list[str] | None = None,
+        descriptions: str | list[str] | None = None,
+        suffixes: str | list[str] | None = None,
+        extensions: str | list[str] | None = [
+            ".con",
+            ".sqd",
+            ".pdf",
+            ".fif",
+            ".ds",
+            ".vhdr",
+            ".set",
+            ".edf",
+            ".bdf",
+            ".EDF",
+            ".snirf",
+            ".cdt",
+            ".mef",
+            ".nwb",
+        ],
+        datatypes: str | list[str] | None = None,
+        check: bool = False,
+    ):
+        if bids_paths is None:
+            bids_paths = mne_bids.find_matching_paths(
+                root=root,
+                subjects=subjects,
+                sessions=sessions,
+                tasks=tasks,
+                acquisitions=acquisitions,
+                runs=runs,
+                processings=processings,
+                recordings=recordings,
+                spaces=spaces,
+                splits=splits,
+                descriptions=descriptions,
+                suffixes=suffixes,
+                extensions=extensions,
+                datatypes=datatypes,
+                check=check,
+                ignore_json=True,
+            )
+            # Filter out _epo.fif files:
+            bids_paths = [
+                bids_path
+                for bids_path in bids_paths
+                if not (bids_path.suffix == "epo" and bids_path.extension == ".fif")
+            ]
+        self.bids_paths = bids_paths
+        self.reader_fn = reader_fn
+        self.pool_size = pool_size
+
+    def __add__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        return BIDSIterableDataset(
+            reader_fn=self.reader_fn,
+            bids_paths=self.bids_paths + other.bids_paths,
+            pool_size=self.pool_size,
+        )
+
+    def __iadd__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        self.bids_paths += other.bids_paths
+        return self
+
+    def __iter__(self):
+        worker_info = get_worker_info()
+        if worker_info is None:  # single-process data loading, return the full iterator
+            bids_paths = self.bids_paths
+        else:  # in a worker process
+            # split workload
+            bids_paths = self.bids_paths[worker_info.id :: worker_info.num_workers]
+
+        pool = []
+        end = False
+        paths_it = iter(random.sample(bids_paths, k=len(bids_paths)))
+        while not (end and len(pool) == 0):
+            while not end and len(pool) < self.pool_size:
+                try:
+                    bids_path = next(paths_it)
+                    ds = self.reader_fn(bids_path)
+                    if ds is None:
+                        print(f"Skipping {bids_path} as it is too short.")
+                        continue
+                    idx = iter(random.sample(range(len(ds)), k=len(ds)))
+                    pool.append((ds, idx))
+                except StopIteration:
+                    end = True
+            i_pool = random.randint(0, len(pool) - 1)
+            ds, idx = pool[i_pool]
+            try:
+                i_ds = next(idx)
+                yield ds[i_ds]
+            except StopIteration:
+                pool.pop(i_pool)

{braindecode-1.2.0.dev175337561 → braindecode-1.2.0.dev176281713}/braindecode/models/__init__.py

@@ -15,7 +15,7 @@ from .eeginception_erp import EEGInceptionERP
 from .eeginception_mi import EEGInceptionMI
 from .eegitnet import EEGITNet
 from .eegminer import EEGMiner
-from .eegnet import EEGNetv1, EEGNetv4
+from .eegnet import EEGNetv4
 from .eegnex import EEGNeX
 from .eegresnet import EEGResNet
 from .eegsimpleconv import EEGSimpleConv
@@ -65,7 +65,6 @@ __all__ = [
     "EEGInceptionMI",
     "EEGITNet",
     "EEGMiner",
-    "EEGNetv1",
     "EEGNetv4",
     "EEGNeX",
     "EEGResNet",

{braindecode-1.2.0.dev175337561 → braindecode-1.2.0.dev176281713}/braindecode/models/attentionbasenet.py

@@ -24,7 +24,7 @@ from braindecode.modules.attention import (
 
 
 class AttentionBaseNet(EEGModuleMixin, nn.Module):
-    """
+    """AttentionBaseNet from Wimpff M et al. (2023) [Martin2023]_.
 
     :bdg-success:`Convolution` :bdg-info:`Small Attention`
 
@@ -57,9 +57,9 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
 
         - *Operations.*
         - **Temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(1, L_t)`` creates a learned
-        FIR-like filter bank with ``n_temporal_filters`` maps.
+          FIR-like filter bank with ``n_temporal_filters`` maps.
         - **Depthwise spatial conv** (:class:`torch.nn.Conv2d`, ``groups=n_temporal_filters``)
-        with kernel ``(n_chans, 1)`` learns per-filter spatial projections over the full montage.
+          with kernel ``(n_chans, 1)`` learns per-filter spatial projections over the full montage.
         - **BatchNorm → ELU → AvgPool → Dropout** stabilize and downsample time.
         - Output shape: ``(B, F2, 1, T₁)`` with ``F2 = n_temporal_filters x spatial_expansion``.
 
@@ -71,8 +71,8 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
 
         - *Operations.*
         - A ``1x1`` conv → BN → activation maps ``F2 → ch_dim`` without changing
-        the temporal length ``T₁`` (shape: ``(B, ch_dim, 1, T₁)``).
-        This sets the embedding width for the attention block.
+          the temporal length ``T₁`` (shape: ``(B, ch_dim, 1, T₁)``).
+          This sets the embedding width for the attention block.
 
     - :class:`_ChannelAttentionBlock` **(temporal refinement + channel attention)**
 
@@ -117,8 +117,10 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
     - **Type.** Channel attention chosen by ``attention_mode`` (SE, ECA, CBAM, CAT, GSoP,
         EncNet, GE, GCT, SRM, CATLite). Most operate purely on channels; CBAM/CAT additionally
         include temporal attention.
+
     - **Shapes.** Input/Output around attention: ``(B, ch_dim, 1, T₁)``. Re-arrangements
         (if any) are internal to the module; the block returns the same shape before pooling.
+
     - **Role.** Re-weights channels (and optionally time) to highlight informative sources
         and suppress distractors, improving SNR ahead of the linear head.
 
@@ -163,10 +165,11 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
     Notes
     -----
     - Sequence length after each stage is computed internally; the final classifier expects
-    a flattened ``ch_dim x T₂`` vector.
+      a flattened ``ch_dim x T₂`` vector.
     - Attention operates on *channel* dimension by design; temporal gating exists only in
-    specific variants (CBAM/CAT).
-
+      specific variants (CBAM/CAT).
+    - The paper and original code with more details about the methodological
+      choices are available at the [Martin2023]_ and [MartinCode]_.
     .. versionadded:: 0.9
 
     Parameters
@@ -195,18 +198,18 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         the depth of the network after the initial layer. Default is 16.
     attention_mode : str, optional
         The type of attention mechanism to apply. If `None`, no attention is applied.
-        - "se" for Squeeze-and-excitation network
-        - "gsop" for Global Second-Order Pooling
-        - "fca" for Frequency Channel Attention Network
-        - "encnet" for context encoding module
-        - "eca" for Efficient channel attention for deep convolutional neural networks
-        - "ge" for Gather-Excite
-        - "gct" for Gated Channel Transformation
-        - "srm" for Style-based Recalibration Module
-        - "cbam" for Convolutional Block Attention Module
-        - "cat" for Learning to collaborate channel and temporal attention
-        from multi-information fusion
-        - "catlite" for Learning to collaborate channel attention
+            - "se" for Squeeze-and-excitation network
+            - "gsop" for Global Second-Order Pooling
+            - "fca" for Frequency Channel Attention Network
+            - "encnet" for context encoding module
+            - "eca" for Efficient channel attention for deep convolutional neural networks
+            - "ge" for Gather-Excite
+            - "gct" for Gated Channel Transformation
+            - "srm" for Style-based Recalibration Module
+            - "cbam" for Convolutional Block Attention Module
+            - "cat" for Learning to collaborate channel and temporal attention
+            from multi-information fusion
+            - "catlite" for Learning to collaborate channel attention
         from multi-information fusion (lite version, cat w/o temporal attention)
     pool_length : int, default=8
         The length of the window for the average pooling operation.

{braindecode-1.2.0.dev175337561 → braindecode-1.2.0.dev176281713}/braindecode/models/eegconformer.py

@@ -27,22 +27,25 @@ class EEGConformer(EEGModuleMixin, nn.Module):
     EEG-Conformer is a *convolution-first* model augmented with a *lightweight transformer
     encoder*. The end-to-end flow is:
 
-    - (i) :class:`_PatchEmbedding` converts the continuous EEG into a compact sequence of tokens via a :class:`ShallowFBCSPNet` temporal–spatial conv stem and temporal pooling;
-    - (ii) :class:`_TransformerEncoder applies small multi-head self-attention to integrate longer-range temporal context across tokens;
+    - (i) :class:`_PatchEmbedding` converts the continuous EEG into a compact sequence of tokens via a
+      :class:`ShallowFBCSPNet` temporal–spatial conv stem and temporal pooling;
+    - (ii) :class:`_TransformerEncoder` applies small multi-head self-attention to integrate
+      longer-range temporal context across tokens;
     - (iii) :class:`_ClassificationHead` aggregates the sequence and performs a linear readout.
-        This preserves the strong inductive biases of shallow CNN filter banks while adding
-        just enough attention to capture dependencies beyond the pooling horizon [song2022]_.
+      This preserves the strong inductive biases of shallow CNN filter banks while adding
+      just enough attention to capture dependencies beyond the pooling horizon [song2022]_.
 
     .. rubric:: Macro Components
 
     - :class:`_PatchEmbedding` **(Shallow conv stem → tokens)**
 
-    - *Operations.*
-        - A temporal convolution (`:class:`torch.nn.Conv2d`) ``(1 x L_t)`` forms a data-driven "filter bank";
-        - A spatial convolution (`:class:`torch.nn.Conv2d`) (n_chans x 1)`` projects across electrodes, collapsing the channel axis into a virtual channel.
-        - **Normalization function** `:class:torch.nn.BatchNorm`
-        - **Activation function** `:class:torch.nn.ELU`
-        - **Average Pooling** `:class:torch.nn.AvgPool` along time (kernel ``(1, P)`` with stride ``(1, S)``)
+        - *Operations.*
+        - A temporal convolution (`:class:torch.nn.Conv2d`) ``(1 x L_t)`` forms a data-driven "filter bank";
+        - A spatial convolution (`:class:torch.nn.Conv2d`) (n_chans x 1)`` projects across electrodes,
+          collapsing the channel axis into a virtual channel.
+        - **Normalization function** :class:`torch.nn.BatchNorm`
+        - **Activation function** :class:`torch.nn.ELU`
+        - **Average Pooling** :class:`torch.nn.AvgPool` along time (kernel ``(1, P)`` with stride ``(1, S)``)
         -  final ``1x1`` :class:`torch.nn.Linear` projection.
 
     The result is rearranged to a token sequence ``(B, S_tokens, D)``, where ``D = n_filters_time``.
@@ -53,7 +56,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
 
     - :class:`_TransformerEncoder` **(context over temporal tokens)**
 
-    - *Operations.*
+        - *Operations.*
         - A stack of ``att_depth`` encoder blocks. :class:`_TransformerEncoderBlock`
         - Each block applies LayerNorm :class:`torch.nn.LayerNorm`
         - Multi-Head Self-Attention (``att_heads``) with dropout + residual :class:`MultiHeadAttention` (:class:`torch.nn.Dropout`)
@@ -67,7 +70,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
 
     - :class:`ClassificationHead` **(aggregation + readout)**
 
-    - *Operations*.
+        - *Operations*.
         - Flatten, :class:`torch.nn.Flatten` the sequence ``(B, S_tokens·D)`` -
         - MLP (:class:`torch.nn.Linear` → activation (default: :class:`torch.nn.ELU`) → :class:`torch.nn.Dropout` → :class:`torch.nn.Linear`)
         - final Linear to classes.

{braindecode-1.2.0.dev175337561 → braindecode-1.2.0.dev176281713}/braindecode/models/eegnet.py

@@ -31,8 +31,7 @@ class EEGNetv4(EEGModuleMixin, nn.Sequential):
 
     .. rubric:: Architectural Overview
 
-    EEGNetv4 is a compact convolutional network designed for EEG decoding with a
-    pipeline that mirrors classical EEG processing:
+    EEGNetv4 is a compact convolutional network designed for EEG decoding with a pipeline that mirrors classical EEG processing:
     - (i) learn temporal frequency-selective filters,
     - (ii) learn spatial filters for those frequencies, and
     - (iii) condense features with depthwise–separable convolutions before a lightweight classifier.
@@ -56,16 +55,16 @@ class EEGNetv4(EEGModuleMixin, nn.Sequential):
 
     .. rubric:: Convolutional Details
 
-    **Temporal.** The initial temporal convs serve as a *learned filter bank*:
-    long 1-D kernels (implemented as 2-D with singleton spatial extent) emphasize oscillatory bands and transients.
-    Because this stage is linear prior to BN/ELU, kernels can be analyzed as FIR filters to reveal each feature’s spectrum [Lawhern2018]_.
+    - **Temporal.** The initial temporal convs serve as a *learned filter bank*:
+      long 1-D kernels (implemented as 2-D with singleton spatial extent) emphasize oscillatory bands and transients.
+      Because this stage is linear prior to BN/ELU, kernels can be analyzed as FIR filters to reveal each feature’s spectrum [Lawhern2018]_.
 
-    **Spatial.** The depthwise spatial conv spans the full channel axis (kernel height = #electrodes; temporal size = 1).
-    With ``groups = F1``, each temporal filter learns its own set of ``D`` spatial projections—akin to CSP, learned end-to-end and
-    typically regularized with max-norm.
+    - **Spatial.** The depthwise spatial conv spans the full channel axis (kernel height = #electrodes; temporal size = 1).
+      With ``groups = F1``, each temporal filter learns its own set of ``D`` spatial projections—akin to CSP, learned end-to-end and
+      typically regularized with max-norm.
 
-    **Spectral.** No explicit Fourier/wavelet transform is used. Frequency structure
-    is captured implicitly by the temporal filter bank; later depthwise temporal kernels act as short-time integrators/refiners.
+    - **Spectral.** No explicit Fourier/wavelet transform is used. Frequency structure
+      is captured implicitly by the temporal filter bank; later depthwise temporal kernels act as short-time integrators/refiners.
 
     .. rubric:: Additional Comments
 

{braindecode-1.2.0.dev175337561 → braindecode-1.2.0.dev176281713}/braindecode/models/eegnex.py

@@ -16,9 +16,124 @@ from braindecode.modules import Conv2dWithConstraint, LinearWithConstraint
 class EEGNeX(EEGModuleMixin, nn.Module):
     """EEGNeX model from Chen et al. (2024) [eegnex]_.
 
+    :bdg-success:`Convolution`
+
     .. figure:: https://braindecode.org/dev/_static/model/eegnex.jpg
         :align: center
         :alt: EEGNeX Architecture
+        :width: 620px
+
+    .. rubric:: Architectural Overview
+
+    EEGNeX is a **purely convolutional** architecture that refines the EEGNet-style stem
+    and deepens the temporal stack with **dilated temporal convolutions**. The end-to-end
+    flow is:
+
+    - (i) **Block-1/2**: two temporal convolutions ``(1 x L)`` with BN refine a
+      learned FIR-like *temporal filter bank* (no pooling yet);
+    - (ii) **Block-3**: depthwise **spatial** convolution across electrodes
+      ``(n_chans x 1)`` with max-norm constraint, followed by ELU → AvgPool (time) → Dropout;
+    - (iii) **Block-4/5**: two additional **temporal** convolutions with increasing **dilation**
+      to expand the receptive field; the last block applies ELU → AvgPool → Dropout → Flatten;
+    - (iv) **Classifier**: a max-norm–constrained linear layer.
+
+    The published work positions EEGNeX as a compact, conv-only alternative that consistently
+    outperforms prior baselines across MOABB-style benchmarks, with the popular
+    “EEGNeX-8,32” shorthand denoting *8 temporal filters* and *kernel length 32*.
+
+
+    .. rubric:: Macro Components
+
+    - **Block-1 / Block-2 — Temporal filter (learned).**
+
+       - *Operations.*
+       - :class:`torch.nn.Conv2d` with kernels ``(1, L)``
+       - :class:`torch.nn.BatchNorm2d` (no nonlinearity until Block-3, mirroring a linear FIR analysis stage).
+         These layers set up frequency-selective detectors before spatial mixing.
+
+    - *Interpretability.* Kernels can be inspected as FIR filters; two stacked temporal
+      convs allow longer effective kernels without parameter blow-up.
+
+    - **Block-3 — Spatial projection + condensation.**
+
+        - *Operations.*
+        - :class:`braindecode.modules.Conv2dWithConstraint` with kernel``(n_chans, 1)``
+          and ``groups = filter_2`` (depthwise across filters)
+        - :class:`torch.nn.BatchNorm2d`
+        - :class:`torch.nn.ELU`
+        - :class:`torch.nn.AvgPool2d` (time)
+        - :class:`torch.nn.Dropout`.
+
+    *Role.* Learns per-filter spatial patterns over the **full montage** while temporal
+      pooling stabilizes and compresses features; max-norm encourages well-behaved spatial
+      weights similar to EEGNet practice.
+
+    - **Block-4 / Block-5 — Dilated temporal integration.**
+
+        - *Operations.*
+        - :class:`torch.nn.Conv2d` with kernels ``(1, k)`` and **dilations**
+        (e.g., 2 then 4);
+        - :class:`torch.nn.BatchNorm2d`
+        - :class:`torch.nn.ELU`
+        - :class:`torch.nn.AvgPool2d` (time)
+        - :class:`torch.nn.Dropout`
+        - :class:`torch.nn.Flatten`.
+
+    *Role.* Expands the temporal receptive field efficiently to capture rhythms and
+      long-range context after condensation.
+
+    - **Final Classifier — Max-norm linear.**
+
+        - *Operations.*
+        - :class:`braindecode.modules.LinearWithConstraint` maps the flattened
+          vector to the target classes; the max-norm constraint regularizes the readout.
+
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+      Blocks 1-2 learn the primary filter bank (oscillations/transients), while Blocks 4-5
+      use **dilation** to integrate over longer horizons without extra pooling. The final
+      AvgPool in Block-5 sets the output token rate and helps noise suppression.
+
+    - **Spatial (how electrodes are processed).**
+      A *single* depthwise spatial conv (Block-3) spans the entire electrode set
+      (kernel ``(n_chans, 1)``), producing per-temporal-filter topographies; no cross-filter
+      mixing occurs at this stage, aiding interpretability.
+
+    - **Spectral (how frequency content is captured).**
+      Frequency selectivity emerges from the learned temporal kernels; dilation broadens effective
+      bandwidth coverage by composing multiple scales.
+
+    .. rubric:: Additional Mechanisms
+
+    - **EEGNeX-8,32 naming.** “8,32” indicates *8 temporal filters* and *kernel length 32*,
+      reflecting the paper's ablation path from EEGNet-8,2 toward thicker temporal kernels
+      and a deeper conv stack.
+    - **Max-norm constraints.** Spatial (Block-3) and final linear layers use max-norm
+      regularization—standard in EEG CNNs—to reduce overfitting and encourage stable spatial
+      patterns.
+
+    .. rubric:: Usage and Configuration
+
+    - **Kernel schedule.** Start with the canonical **EEGNeX-8,32** (``filter_1=8``,
+      ``kernel_block_1_2=32``) and keep **Block-3** depth multiplier modest (e.g., 2) to match
+      the paper's “pure conv” profile.
+    - **Pooling vs. dilation.** Use pooling in Blocks 3 and 5 to control compute and variance;
+      increase dilations (Blocks 4-5) to widen temporal context when windows are short.
+    - **Regularization.** Combine dropout (Blocks 3 & 5) with max-norm on spatial and
+      classifier layers; prefer ELU activations for stable training on small EEG datasets.
+
+
+    Notes
+    -----
+    - The braindecode implementation follows the paper's conv-only design with five blocks
+      and reproduces the depthwise spatial step and dilated temporal stack. See the class
+      reference for exact kernel sizes, dilations, and pooling defaults. You can check the
+      original implementation at [EEGNexCode]_.
+
+    .. versionadded:: 1.1
+
 
     Parameters
     ----------
@@ -45,12 +160,6 @@ class EEGNeX(EEGModuleMixin, nn.Module):
     avg_pool_block5 : tuple[int, int], optional
         Pooling size for block 5. Default is (1, 8).
 
-    Notes
-    -----
-    This implementation is not guaranteed to be correct, has not been checked
-    by original authors, only reimplemented from the paper description and
-    source code in tensorflow [EEGNexCode]_.
-
     References
     ----------
     .. [eegnex] Chen, X., Teng, X., Chen, H., Pan, Y., & Geyer, P. (2024).