braindecode 1.0.0__tar.gz → 1.1.0__tar.gz

{braindecode-1.0.0/braindecode.egg-info → braindecode-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.0.0
+Version: 1.1.0
 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>
@@ -29,12 +29,12 @@ Requires-Dist: pandas
 Requires-Dist: scipy
 Requires-Dist: matplotlib
 Requires-Dist: h5py
-Requires-Dist: skorch
-Requires-Dist: torch
-Requires-Dist: torchaudio
+Requires-Dist: skorch~=0.14.0
+Requires-Dist: torch~=2.0
+Requires-Dist: torchaudio~=2.0
 Requires-Dist: einops
 Requires-Dist: joblib
-Requires-Dist: torchinfo
+Requires-Dist: torchinfo~=1.8
 Requires-Dist: wfdb
 Requires-Dist: h5py
 Requires-Dist: linear_attention_transformer

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/augmentation/transforms.py

@@ -6,7 +6,6 @@
 
 import warnings
 from numbers import Real
-from typing import Callable
 
 import numpy as np
 import torch

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/datautil/__init__.py

@@ -47,3 +47,6 @@ def __getattr__(name):
         return windowers.__dict__[name]
 
     raise AttributeError("No possible import named " + name)
+
+
+__all__ = ["load_concat_dataset", "save_concat_dataset", "_check_save_dir_empty"]

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/datautil/serialization.py

@@ -107,7 +107,14 @@ def _outdated_load_concat_dataset(path, preload, ids_to_load=None, target_name=N
 def _load_signals_and_description(path, preload, is_raw, ids_to_load=None):
     all_signals = []
     file_name = "{}-raw.fif" if is_raw else "{}-epo.fif"
-    description_df = pd.read_json(path / "description.json")
+    description_df = pd.read_json(
+        path / "description.json", typ="series", convert_dates=False
+    )
+
+    if "timestamp" in description_df.index:
+        timestamp_numeric = pd.to_numeric(description_df["timestamp"])
+        description_df["timestamp"] = pd.to_datetime(timestamp_numeric)
+
     if ids_to_load is None:
         file_names = path.glob(f"*{file_name.lstrip('{}')}")
         # Extract ids, e.g.,
@@ -242,7 +249,11 @@ def _load_parallel(path, i, preload, is_raw, has_stored_windows):
     signals = _load_signals(fif_file_path, preload, is_raw)
 
     description_file_path = sub_dir / "description.json"
-    description = pd.read_json(description_file_path, typ="series")
+    description = pd.read_json(description_file_path, typ="series", convert_dates=False)
+
+    # if 'timestamp' in description.index:
+    #     timestamp_numeric = pd.to_numeric(description['timestamp'])
+    #     description['timestamp'] = pd.to_datetime(timestamp_numeric, unit='s')
 
     target_file_path = sub_dir / "target_name.json"
     target_name = None

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/functional/__init__.py

@@ -8,3 +8,15 @@ from .functions import (
     square,
 )
 from .initialization import glorot_weight_zero_bias, rescale_parameter
+
+__all__ = [
+    "_get_gaussian_kernel1d",
+    "drop_path",
+    "hilbert_freq",
+    "identity",
+    "plv_time",
+    "safe_log",
+    "square",
+    "glorot_weight_zero_bias",
+    "rescale_parameter",
+]

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/functional/functions.py

@@ -1,7 +1,6 @@
 # Authors: Robin Schirrmeister <robintibor@gmail.com>
 #
 # License: BSD (3-clause)
-import math
 
 import torch
 import torch.nn.functional as F

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/models/__init__.py

@@ -50,3 +50,51 @@ from .util import _init_models_dict, models_mandatory_parameters
 # Call this last in order to make sure the dataset list is populated with
 # the models imported in this file.
 _init_models_dict()
+
+__all__ = [
+    "ATCNet",
+    "AttentionBaseNet",
+    "EEGModuleMixin",
+    "BIOT",
+    "ContraWR",
+    "CTNet",
+    "Deep4Net",
+    "DeepSleepNet",
+    "EEGConformer",
+    "EEGInceptionERP",
+    "EEGInceptionMI",
+    "EEGITNet",
+    "EEGMiner",
+    "EEGNetv1",
+    "EEGNetv4",
+    "EEGNeX",
+    "EEGResNet",
+    "EEGSimpleConv",
+    "EEGTCNet",
+    "FBCNet",
+    "FBLightConvNet",
+    "FBMSNet",
+    "HybridNet",
+    "IFNet",
+    "Labram",
+    "MSVTNet",
+    "SCCNet",
+    "ShallowFBCSPNet",
+    "SignalJEPA",
+    "SignalJEPA_Contextual",
+    "SignalJEPA_PostLocal",
+    "SignalJEPA_PreLocal",
+    "SincShallowNet",
+    "SleepStagerBlanco2020",
+    "SleepStagerChambon2018",
+    "SleepStagerEldele2021",
+    "SPARCNet",
+    "SyncNet",
+    "BDTCN",
+    "TCN",
+    "TIDNet",
+    "TSceptionV1",
+    "USleep",
+    "_init_models_dict",
+    "models_mandatory_parameters",
+]

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/models/atcnet.py

@@ -5,6 +5,7 @@ import math
 
 import torch
 from einops.layers.torch import Rearrange
+from mne.utils import warn
 from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
@@ -69,9 +70,6 @@ class ATCNet(EEGModuleMixin, nn.Module):
     tcn_kernel_size : int
         Temporal kernel size used in TCN block, denoted Kt in table 1 of the
         paper [1]_. Defaults to 4 as in [1]_.
-    tcn_n_filters : int
-        Number of filters used in TCN convolutional layers (Ft). Defaults to
-        32 as in [1]_.
     tcn_dropout : float
         Dropout probability used in the TCN block, denoted pt in table 1
         of the paper [1]_. Defaults to 0.3 as in [1]_.
@@ -117,7 +115,6 @@ class ATCNet(EEGModuleMixin, nn.Module):
         att_drop_prob=0.5,
         tcn_depth=2,
         tcn_kernel_size=4,
-        tcn_n_filters=32,
         tcn_drop_prob=0.3,
         tcn_activation: nn.Module = nn.ELU,
         concat=False,
@@ -134,6 +131,45 @@ class ATCNet(EEGModuleMixin, nn.Module):
             sfreq=sfreq,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+
+        # Validate and adjust parameters based on input size
+
+        min_len_tcn = (tcn_kernel_size - 1) * (2 ** (tcn_depth - 1)) + 1
+        # Minimum length required to get at least one sliding window
+        min_len_sliding = n_windows + min_len_tcn - 1
+        # Minimum input size that produces the required feature map length
+        min_n_times = min_len_sliding * conv_block_pool_size_1 * conv_block_pool_size_2
+
+        # 2. If the input is shorter, calculate a scaling factor
+        if self.n_times < min_n_times:
+            scaling_factor = self.n_times / min_n_times
+            warn(
+                f"n_times ({self.n_times}) is smaller than the minimum required "
+                f"({min_n_times}) for the current model parameters configuration. "
+                "Adjusting parameters to ensure compatibility."
+                "Reducing the kernel, pooling, and stride sizes accordingly."
+                "Scaling factor: {:.2f}".format(scaling_factor),
+                UserWarning,
+            )
+            conv_block_kernel_length_1 = max(
+                1, int(conv_block_kernel_length_1 * scaling_factor)
+            )
+            conv_block_kernel_length_2 = max(
+                1, int(conv_block_kernel_length_2 * scaling_factor)
+            )
+            conv_block_pool_size_1 = max(
+                1, int(conv_block_pool_size_1 * scaling_factor)
+            )
+            conv_block_pool_size_2 = max(
+                1, int(conv_block_pool_size_2 * scaling_factor)
+            )
+
+            # n_windows should be at least 1
+            n_windows = max(1, int(n_windows * scaling_factor))
+
+            # tcn_kernel_size must be at least 2 for dilation to work
+            tcn_kernel_size = max(2, int(tcn_kernel_size * scaling_factor))
+
         self.conv_block_n_filters = conv_block_n_filters
         self.conv_block_kernel_length_1 = conv_block_kernel_length_1
         self.conv_block_kernel_length_2 = conv_block_kernel_length_2
@@ -147,12 +183,11 @@ class ATCNet(EEGModuleMixin, nn.Module):
         self.att_dropout = att_drop_prob
         self.tcn_depth = tcn_depth
         self.tcn_kernel_size = tcn_kernel_size
-        self.tcn_n_filters = tcn_n_filters
         self.tcn_dropout = tcn_drop_prob
         self.tcn_activation = tcn_activation
         self.concat = concat
         self.max_norm_const = max_norm_const
-
+        self.tcn_n_filters = int(self.conv_block_depth_mult * self.conv_block_n_filters)
         map = dict()
         for w in range(self.n_windows):
             map[f"max_norm_linears.[{w}].weight"] = f"final_layer.[{w}].weight"
@@ -197,13 +232,13 @@ class ATCNet(EEGModuleMixin, nn.Module):
                     *[
                         _TCNResidualBlock(
                             in_channels=self.F2,
-                            kernel_size=tcn_kernel_size,
-                            n_filters=tcn_n_filters,
-                            dropout=tcn_drop_prob,
-                            activation=tcn_activation,
+                            kernel_size=self.tcn_kernel_size,
+                            n_filters=self.tcn_n_filters,
+                            dropout=self.tcn_dropout,
+                            activation=self.tcn_activation,
                             dilation=2**i,
                         )
-                        for i in range(tcn_depth)
+                        for i in range(self.tcn_depth)
                     ]
                 )
                 for _ in range(self.n_windows)

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/models/attentionbasenet.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import math
 
 from einops.layers.torch import Rearrange
+from mne.utils import warn
 from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
@@ -162,6 +163,33 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         )
         del n_outputs, n_chans, chs_info, n_times, sfreq, input_window_seconds
 
+        min_n_times_required = self._get_min_n_times(
+            pool_length_inp,
+            pool_stride_inp,
+            pool_length,
+        )
+
+        if self.n_times < min_n_times_required:
+            scaling_factor = self.n_times / min_n_times_required
+            warn(
+                f"n_times ({self.n_times}) is smaller than the minimum required "
+                f"({min_n_times_required}) for the current model parameters configuration. "
+                "Adjusting parameters to ensure compatibility."
+                "Reducing the kernel, pooling, and stride sizes accordingly.\n"
+                "Scaling factor: {:.2f}".format(scaling_factor),
+                UserWarning,
+            )
+            # 3. Scale down all temporal parameters proportionally
+            # Use max(1, ...) to ensure parameters remain valid
+            temp_filter_length_inp = max(
+                1, int(temp_filter_length_inp * scaling_factor)
+            )
+            pool_length_inp = max(1, int(pool_length_inp * scaling_factor))
+            pool_stride_inp = max(1, int(pool_stride_inp * scaling_factor))
+            temp_filter_length = max(1, int(temp_filter_length * scaling_factor))
+            pool_length = max(1, int(pool_length * scaling_factor))
+            pool_stride = max(1, int(pool_stride * scaling_factor))
+
         self.input_block = _FeatureExtractor(
             n_chans=self.n_chans,
             n_temporal_filters=n_temporal_filters,
@@ -231,6 +259,27 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
             seq_lengths.append(int(out))
         return seq_lengths
 
+    @staticmethod
+    def _get_min_n_times(
+        pool_length_inp: int,
+        pool_stride_inp: int,
+        pool_length: int,
+    ) -> int:
+        """
+        Calculates the minimum n_times required for the model to work
+        with the given parameters.
+
+        The calculation is based on reversing the pooling operations to
+        ensure the input to each is valid.
+        """
+        # The input to the second pooling layer must be at least its kernel size.
+        min_len_for_second_pool = pool_length
+
+        # Reverse the first pooling operation to find the required input size.
+        # Formula: min_L_in = Stride * (min_L_out - 1) + Kernel
+        min_len = pool_stride_inp * (min_len_for_second_pool - 1) + pool_length_inp
+        return min_len
+
 
 class _FeatureExtractor(nn.Module):
     """

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/models/biot.py

@@ -87,6 +87,10 @@ class BIOT(EEGModuleMixin, nn.Module):
         input_window_seconds=None,
         activation: nn.Module = nn.ELU,
         drop_prob: float = 0.5,
+        # Parameters for the encoder
+        max_seq_len: int = 1024,
+        attn_dropout=0.2,
+        attn_layer_dropout=0.2,
     ):
         super().__init__(
             n_outputs=n_outputs,
@@ -123,14 +127,29 @@ class BIOT(EEGModuleMixin, nn.Module):
                 UserWarning,
             )
             hop_length = self.sfreq // 2
+
+        if self.input_window_seconds < 1.0:
+            warning_msg = (
+                "The input window is less than 1 second, which may not be "
+                "sufficient for the model to learn meaningful representations."
+                "Changing the `n_fft` to `n_times`."
+            )
+            warn(warning_msg, UserWarning)
+            self.n_fft = self.n_times
+        else:
+            self.n_fft = int(self.sfreq)
+
         self.encoder = _BIOTEncoder(
             emb_size=emb_size,
             att_num_heads=att_num_heads,
             n_layers=n_layers,
             n_chans=self.n_chans,
-            n_fft=self.sfreq,
+            n_fft=self.n_fft,
             hop_length=hop_length,
             drop_prob=drop_prob,
+            max_seq_len=max_seq_len,
+            attn_dropout=attn_dropout,
+            attn_layer_dropout=attn_layer_dropout,
         )
 
         self.final_layer = _ClassificationHead(
@@ -231,12 +250,11 @@ class _ClassificationHead(nn.Sequential):
 
     def __init__(self, emb_size: int, n_outputs: int, activation: nn.Module = nn.ELU):
         super().__init__()
-        self.classification_head = nn.Sequential(
-            activation(),
-            nn.Linear(emb_size, n_outputs),
-        )
+        self.activation_layer = activation()
+        self.classification_head = nn.Linear(emb_size, n_outputs)
 
     def forward(self, x):
+        x = self.activation_layer(x)
         out = self.classification_head(x)
         return out
 
@@ -344,6 +362,9 @@ class _BIOTEncoder(nn.Module):
         n_fft=200,  # Related with the frequency resolution
         hop_length=100,
         drop_prob: float = 0.1,
+        max_seq_len: int = 1024,  # The maximum sequence length
+        attn_dropout=0.2,  # dropout post-attention
+        attn_layer_dropout=0.2,  # dropout right after self-attention layer
     ):
         super().__init__()
 
@@ -357,9 +378,9 @@ class _BIOTEncoder(nn.Module):
             dim=emb_size,
             heads=att_num_heads,
             depth=n_layers,
-            max_seq_len=1024,
-            attn_layer_dropout=0.2,  # dropout right after self-attention layer
-            attn_dropout=0.2,  # dropout post-attention
+            max_seq_len=max_seq_len,
+            attn_layer_dropout=attn_layer_dropout,
+            attn_dropout=attn_dropout,
         )
         self.positional_encoding = _PositionalEncoding(emb_size, drop_prob=drop_prob)
 

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/models/contrawr.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import torch
 import torch.nn as nn
+from mne.utils import warn
 
 from braindecode.models.base import EEGModuleMixin
 
@@ -57,6 +58,9 @@ class ContraWR(EEGModuleMixin, nn.Module):
         steps=20,
         activation: nn.Module = nn.ELU,
         drop_prob: float = 0.5,
+        stride_res: int = 2,
+        kernel_size_res: int = 3,
+        padding_res: int = 1,
         # Another way to pass the EEG parameters
         chs_info=None,
         n_times=None,
@@ -74,7 +78,17 @@ class ContraWR(EEGModuleMixin, nn.Module):
         if not isinstance(res_channels, list):
             raise ValueError("res_channels must be a list of integers.")
 
-        self.n_fft = int(self.sfreq)
+        if self.input_window_seconds < 1.0:
+            warning_msg = (
+                "The input window is less than 1 second, which may not be "
+                "sufficient for the model to learn meaningful representations."
+                "changing the `n_fft` to `n_times`."
+            )
+            warn(warning_msg, UserWarning)
+            self.n_fft = self.n_times
+        else:
+            self.n_fft = int(self.sfreq)
+
         self.steps = steps
 
         res_channels = [self.n_chans] + res_channels + [emb_size]
@@ -89,19 +103,22 @@ class ContraWR(EEGModuleMixin, nn.Module):
                 _ResBlock(
                     in_channels=res_channels[i],
                     out_channels=res_channels[i + 1],
-                    stride=2,
+                    stride=stride_res,
                     use_downsampling=True,
                     pooling=True,
                     drop_prob=drop_prob,
+                    kernel_size=kernel_size_res,
+                    padding=padding_res,
+                    activation=activation,
                 )
                 for i in range(len(res_channels) - 1)
             ]
         )
+        self.adaptative_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.flatten_layer = nn.Flatten()
 
-        self.final_layer = nn.Sequential(
-            activation(),
-            nn.Linear(emb_size, self.n_outputs),
-        )
+        self.activation_layer = activation()
+        self.final_layer = nn.Linear(emb_size, self.n_outputs)
 
     def forward(self, X: torch.Tensor) -> torch.Tensor:
         """
@@ -118,9 +135,13 @@ class ContraWR(EEGModuleMixin, nn.Module):
         """
         X = self.torch_stft(X)
 
-        for conv in self.convs[:-1]:
+        for conv in self.convs:
             X = conv.forward(X)
-        emb = self.convs[-1](X).squeeze(-1).squeeze(-1)
+
+        emb = self.adaptative_pool(X)
+        emb = self.flatten_layer(emb)
+        emb = self.activation_layer(emb)
+
         return self.final_layer(emb)
 
 

{braindecode-1.0.0 → braindecode-1.1.0}/braindecode/models/ctnet.py

@@ -10,6 +10,7 @@ classification from Wei Zhao et al. (2024).
 from __future__ import annotations
 
 import math
+from typing import Optional
 
 import torch
 from einops.layers.torch import Rearrange
@@ -57,7 +58,7 @@ class CTNet(EEGModuleMixin, nn.Module):
         Activation function to use in the network.
     heads : int, default=4
         Number of attention heads in the Transformer encoder.
-    emb_size : int, default=40
+    emb_size : int or None, default=None
         Embedding size (dimensionality) for the Transformer encoder.
     depth : int, default=6
         Number of encoder layers in the Transformer.
@@ -110,11 +111,11 @@ class CTNet(EEGModuleMixin, nn.Module):
         drop_prob_final: float = 0.5,
         # other parameters
         heads: int = 4,
-        emb_size: int = 40,
+        emb_size: Optional[int] = 40,
         depth: int = 6,
-        n_filters_time: int = 20,
+        n_filters_time: Optional[int] = None,
         kernel_size: int = 64,
-        depth_multiplier: int = 2,
+        depth_multiplier: Optional[int] = 2,
         pool_size_1: int = 8,
         pool_size_2: int = 8,
     ):
@@ -128,21 +129,18 @@ class CTNet(EEGModuleMixin, nn.Module):
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
 
-        self.emb_size = emb_size
         self.activation_patch = activation_patch
         self.activation_transformer = activation_transformer
-
-        self.n_filters_time = n_filters_time
         self.drop_prob_cnn = drop_prob_cnn
         self.pool_size_1 = pool_size_1
         self.pool_size_2 = pool_size_2
-        self.depth_multiplier = depth_multiplier
         self.kernel_size = kernel_size
         self.drop_prob_posi = drop_prob_posi
         self.drop_prob_final = drop_prob_final
-
+        self.heads = heads
+        self.depth = depth
         # n_times - pool_size_1 / p
-        sequence_length = math.floor(
+        self.sequence_length = math.floor(
             (
                 math.floor((self.n_times - self.pool_size_1) / self.pool_size_1 + 1)
                 - self.pool_size_2
@@ -151,6 +149,10 @@ class CTNet(EEGModuleMixin, nn.Module):
             + 1
         )
 
+        self.depth_multiplier, self.n_filters_time, self.emb_size = self._resolve_dims(
+            depth_multiplier, n_filters_time, emb_size
+        )
+
         # Layers
         self.ensuredim = Rearrange("batch nchans time -> batch 1 nchans time")
         self.flatten = nn.Flatten()
@@ -167,14 +169,17 @@ class CTNet(EEGModuleMixin, nn.Module):
         )
 
         self.position = _PositionalEncoding(
-            emb_size=emb_size,
+            emb_size=self.emb_size,
             drop_prob=self.drop_prob_posi,
             n_times=self.n_times,
             pool_size=self.pool_size_1,
         )
 
         self.trans = _TransformerEncoder(
-            heads, depth, emb_size, activation=self.activation_transformer
+            self.heads,
+            self.depth,
+            self.emb_size,
+            activation=self.activation_transformer,
         )
 
         self.flatten_drop_layer = nn.Sequential(
@@ -183,7 +188,8 @@ class CTNet(EEGModuleMixin, nn.Module):
         )
 
         self.final_layer = nn.Linear(
-            in_features=emb_size * sequence_length, out_features=self.n_outputs
+            in_features=int(self.emb_size * self.sequence_length),
+            out_features=self.n_outputs,
         )
 
     def forward(self, x: Tensor) -> Tensor:
@@ -210,6 +216,86 @@ class CTNet(EEGModuleMixin, nn.Module):
         out = self.final_layer(flatten_feature)
         return out
 
+    @staticmethod
+    def _resolve_dims(
+        depth_multiplier: Optional[int],
+        n_filters_time: Optional[int],
+        emb_size: Optional[int],
+    ) -> tuple[int, int, int]:
+        # Basic type/positivity checks for provided values
+        for name, val in (
+            ("depth_multiplier", depth_multiplier),
+            ("n_filters_time", n_filters_time),
+            ("emb_size", emb_size),
+        ):
+            if val is not None:
+                if not isinstance(val, int):
+                    raise TypeError(f"{name} must be int, got {type(val).__name__}")
+                if val <= 0:
+                    raise ValueError(f"{name} must be > 0, got {val}")
+
+        missing = [
+            k
+            for k, v in {
+                "depth_multiplier": depth_multiplier,
+                "n_filters_time": n_filters_time,
+                "emb_size": emb_size,
+            }.items()
+            if v is None
+        ]
+
+        if len(missing) >= 2:
+            # Too many unknowns → ambiguous
+            raise ValueError(
+                "Specify exactly two of {depth_multiplier, n_filters_time, emb_size}; the third will be inferred."
+            )
+
+        if len(missing) == 1:
+            # Infer the missing one
+            if missing[0] == "emb_size":
+                assert depth_multiplier is not None and n_filters_time is not None
+                emb_size = depth_multiplier * n_filters_time
+            elif missing[0] == "n_filters_time":
+                assert emb_size is not None and depth_multiplier is not None
+                if emb_size % depth_multiplier != 0:
+                    raise ValueError(
+                        f"emb_size={emb_size} must be divisible by depth_multiplier={depth_multiplier}"
+                    )
+                n_filters_time = emb_size // depth_multiplier
+            else:  # missing depth_multiplier
+                assert emb_size is not None and n_filters_time is not None
+                if emb_size % n_filters_time != 0:
+                    raise ValueError(
+                        f"emb_size={emb_size} must be divisible by n_filters_time={n_filters_time}"
+                    )
+                depth_multiplier = emb_size // n_filters_time
+
+        else:
+            # All provided: enforce consistency
+            assert (
+                depth_multiplier is not None
+                and n_filters_time is not None
+                and emb_size is not None
+            )
+            prod = depth_multiplier * n_filters_time
+            if prod != emb_size:
+                raise ValueError(
+                    "`depth_multiplier * n_filters_time` must equal `emb_size`, "
+                    f"but got {depth_multiplier} * {n_filters_time} = {prod} != {emb_size}. "
+                    "Fix by setting one of: "
+                    f"emb_size={prod}, "
+                    f"n_filters_time={emb_size // depth_multiplier if emb_size % depth_multiplier == 0 else 'not integer'}, "
+                    f"depth_multiplier={emb_size // n_filters_time if emb_size % n_filters_time == 0 else 'not integer'}."
+                )
+
+        # Ensure plain ints for the return type
+        assert (
+            depth_multiplier is not None
+            and n_filters_time is not None
+            and emb_size is not None
+        )
+        return depth_multiplier, n_filters_time, emb_size
+
 
 class _PatchEmbeddingEEGNet(nn.Module):
     def __init__(