braindecode 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

braindecode/models/deep4.py

@@ -3,6 +3,7 @@
 # License: BSD (3-clause)
 
 from einops.layers.torch import Rearrange
+from mne.utils import warn
 from torch import nn
 from torch.nn import init
 
@@ -115,6 +116,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         batch_norm=True,
         batch_norm_alpha=0.1,
         stride_before_pool=False,
+        # Braindecode EEGModuleMixin parameters
         chs_info=None,
         input_window_seconds=None,
         sfreq=None,
@@ -155,6 +157,27 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         self.batch_norm_alpha = batch_norm_alpha
         self.stride_before_pool = stride_before_pool
 
+        min_n_times = self._get_min_n_times()
+        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,
+            )
+            # Calculate a scaling factor to adjust temporal parameters
+            # Apply the scaling factor to all temporal kernel and pooling sizes
+            self.filter_time_length = max(
+                1, int(self.filter_time_length * scaling_factor)
+            )
+            self.pool_time_length = max(1, int(self.pool_time_length * scaling_factor))
+            self.pool_time_stride = max(1, int(self.pool_time_stride * scaling_factor))
+            self.filter_length_2 = max(1, int(self.filter_length_2 * scaling_factor))
+            self.filter_length_3 = max(1, int(self.filter_length_3 * scaling_factor))
+            self.filter_length_4 = max(1, int(self.filter_length_4 * scaling_factor))
         # For the load_state_dict
         # When padronize all layers,
         # add the old's parameters here
@@ -268,7 +291,6 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
             self, self.n_filters_3, self.n_filters_4, self.filter_length_4, 4
         )
 
-        # self.add_module('drop_classifier', nn.Dropout(p=self.drop_prob))
         self.eval()
         if self.final_conv_length == "auto":
             self.final_conv_length = self.get_output_shape()[2]
@@ -299,7 +321,7 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         if self.split_first_layer:
             init.xavier_uniform_(self.conv_time_spat.conv_spat.weight, gain=1)
             if not self.batch_norm:
-                init.constant_(self.conv_spat.bias, 0)
+                init.constant_(self.conv_time_spat.conv_spat.bias, 0)
         if self.batch_norm:
             init.constant_(self.bnorm.weight, 1)
             init.constant_(self.bnorm.bias, 0)
@@ -320,3 +342,31 @@ class Deep4Net(EEGModuleMixin, nn.Sequential):
         init.constant_(self.final_layer.conv_classifier.bias, 0)
 
         self.train()
+
+    def _get_min_n_times(self) -> int:
+        """
+        Calculate the minimum number of time samples required for the model
+        to work with the given temporal parameters.
+        """
+        # Start with the minimum valid output length of the network (1)
+        min_len = 1
+
+        # List of conv kernel sizes and pool parameters for the 4 blocks, in reverse order
+        # Each tuple: (filter_length, pool_length, pool_stride)
+        block_params = [
+            (self.filter_length_4, self.pool_time_length, self.pool_time_stride),
+            (self.filter_length_3, self.pool_time_length, self.pool_time_stride),
+            (self.filter_length_2, self.pool_time_length, self.pool_time_stride),
+            (self.filter_time_length, self.pool_time_length, self.pool_time_stride),
+        ]
+
+        # Work backward from the last layer to the input
+        for filter_len, pool_len, pool_stride in block_params:
+            # Reverse the pooling operation
+            # L_in = stride * (L_out - 1) + kernel_size
+            min_len = pool_stride * (min_len - 1) + pool_len
+            # Reverse the convolution operation (assuming stride=1)
+            # L_in = L_out + kernel_size - 1
+            min_len = min_len + filter_len - 1
+
+        return min_len

braindecode/models/eegconformer.py

@@ -2,11 +2,8 @@
 #
 # License: BSD (3-clause)
 import warnings
-from typing import Optional
 
 import torch
-import torch.nn.functional as F
-from einops import rearrange
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
@@ -150,6 +147,8 @@ class EEGConformer(EEGModuleMixin, nn.Module):
         if final_fc_length == "auto":
             assert self.n_times is not None
             self.final_fc_length = self.get_fc_size()
+        else:
+            self.final_fc_length = final_fc_length
 
         self.transformer = _TransformerEncoder(
             att_depth=att_depth,

braindecode/models/eeginception_mi.py

@@ -66,9 +66,9 @@ class EEGInceptionMI(EEGModuleMixin, nn.Module):
         n_outputs=None,
         input_window_seconds=None,
         sfreq=250,
-        n_convs=5,
-        n_filters=48,
-        kernel_unit_s=0.1,
+        n_convs: int = 5,
+        n_filters: int = 48,
+        kernel_unit_s: float = 0.1,
         activation: nn.Module = nn.ReLU,
         chs_info=None,
         n_times=None,
@@ -307,6 +307,12 @@ class _InceptionModuleMI(nn.Module):
 
         X2 = self.pooling(X)
         X2 = self.pooling_conv(X2)
+        # Get the target length from one of the conv branches
+        target_len = X1[0].shape[-1]
+
+        # Crop the pooling output if its length does not match
+        if X2.shape[-1] != target_len:
+            X2 = X2[..., :target_len]
 
         out = torch.cat(X1 + [X2], 1)
 

braindecode/models/eegitnet.py

@@ -1,7 +1,6 @@
 # Authors: Ghaith Bouallegue <ghaithbouallegue@gmail.com>
 #
 # License: BSD-3
-import torch
 from einops.layers.torch import Rearrange
 from torch import nn
 

braindecode/models/eegminer.py

@@ -10,7 +10,6 @@ from functools import partial
 import torch
 from einops.layers.torch import Rearrange
 from torch import nn
-from torch.fft import fftfreq
 
 import braindecode.functional as F
 from braindecode.models.base import EEGModuleMixin

braindecode/models/eegnet.py

@@ -14,7 +14,6 @@ from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import (
     Conv2dWithConstraint,
     Ensure4d,
-    Expression,
     LinearWithConstraint,
     SqueezeFinalOutput,
 )

braindecode/models/fbcnet.py

@@ -5,7 +5,7 @@ from typing import Any
 import torch
 from einops.layers.torch import Rearrange
 from mne.utils import warn
-from torch import Tensor, nn
+from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import (

braindecode/models/fbmsnet.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from functools import partial
 from typing import Optional, Sequence
 
 import torch

braindecode/models/labram.py

@@ -166,11 +166,23 @@ class Labram(EEGModuleMixin, nn.Module):
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
 
         self.patch_size = patch_size
-        self.n_path = self.n_times // patch_size
         self.num_features = self.emb_size = emb_size
         self.neural_tokenizer = neural_tokenizer
         self.init_scale = init_scale
 
+        if patch_size > self.n_times:
+            warn(
+                f"patch_size ({patch_size}) > n_times ({self.n_times}); "
+                f"setting patch_size = {self.n_times}.",
+                UserWarning,
+            )
+            self.patch_size = self.n_times
+            self.num_features = None
+            self.emb_size = None
+        else:
+            self.patch_size = patch_size
+        self.n_path = self.n_times // self.patch_size
+
         if neural_tokenizer and in_channels != 1:
             warn(
                 "The model is in Neural Tokenizer mode, but the variable "
@@ -220,8 +232,17 @@ class Labram(EEGModuleMixin, nn.Module):
                     emb_dim=self.emb_size,
                 ),
             )
+
+        with torch.no_grad():
+            dummy = torch.zeros(1, self.n_chans, self.n_times)
+            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
+        self.emb_size = out.shape[-1]
+        self.num_features = self.emb_size
+
         # Defining the parameters
-        # Creating a parameter list with cls token
+        # Creating a parameter list with cls token]
         self.cls_token = nn.Parameter(torch.zeros(1, 1, self.emb_size))
         # Positional embedding and time embedding are complementary
         # one is for the spatial information and the other is for the temporal
@@ -366,7 +387,6 @@ class Labram(EEGModuleMixin, nn.Module):
         x : torch.Tensor
             The output of the model.
         """
-
         if self.neural_tokenizer:
             batch_size, nch, n_patch, temporal = self.patch_embed.segment_patch(x).shape
         else:

braindecode/models/msvtnet.py

@@ -1,7 +1,7 @@
 # Authors: Tao Yang <sheeptao@outlook.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com> (braindecode adaptation)
 #
-from typing import Dict, Optional, Type, Union
+from typing import Type, Union
 
 import torch
 import torch.nn as nn

braindecode/models/sccnet.py

@@ -4,6 +4,7 @@
 # License: BSD (3-clause)
 
 import math
+from warnings import warn
 
 import torch
 from einops.layers.torch import Rearrange
@@ -98,9 +99,33 @@ class SCCNet(EEGModuleMixin, nn.Module):
         self.n_spatial_filters_smooth = n_spatial_filters_smooth
         self.drop_prob = drop_prob
 
-        self.samples_100ms = int(math.floor(self.sfreq * 0.1))
-        self.kernel_size_pool = int(self.sfreq * 0.5)
-        # Equivalent to 0.5 seconds
+        # Original logical for SCCNet
+        conv_kernel_time = 0.1  # 100ms
+        pool_kernel_time = 0.5  # 500ms
+
+        # Calculate sample-based sizes from time durations
+        conv_kernel_samples = int(math.floor(self.sfreq * conv_kernel_time))
+        pool_kernel_samples = int(math.floor(self.sfreq * pool_kernel_time))
+
+        # If the input window is too short for the default kernel sizes,
+        # scale them down proportionally.
+        total_kernel_samples = conv_kernel_samples + pool_kernel_samples
+
+        if self.n_times < total_kernel_samples:
+            warning_msg = (
+                f"Input window seconds ({self.input_window_seconds:.2f}s) is smaller than the "
+                f"model's combined kernel sizes ({(total_kernel_samples / self.sfreq):.2f}s). "
+                "Scaling temporal parameters down proportionally."
+            )
+            warn(warning_msg, UserWarning, stacklevel=2)
+
+            scaling_factor = self.n_times / total_kernel_samples
+            conv_kernel_samples = int(math.floor(conv_kernel_samples * scaling_factor))
+            pool_kernel_samples = int(math.floor(pool_kernel_samples * scaling_factor))
+
+        # Ensure kernels are at least 1 sample wide
+        self.samples_100ms = max(1, conv_kernel_samples)
+        self.kernel_size_pool = max(1, pool_kernel_samples)
 
         num_features = self._calc_num_features()
 
@@ -135,7 +160,7 @@ class SCCNet(EEGModuleMixin, nn.Module):
 
         self.dropout = nn.Dropout(self.drop_prob)
         self.temporal_smoothing = nn.AvgPool2d(
-            kernel_size=(1, int(self.sfreq / 2)),
+            kernel_size=(1, self.kernel_size_pool),
             stride=(1, self.samples_100ms),
         )
 

braindecode/models/signal_jepa.py

@@ -8,7 +8,6 @@ from copy import deepcopy
 from typing import Any, Sequence
 
 import torch
-from einops import parse_shape, rearrange, repeat
 from einops.layers.torch import Rearrange
 from torch import nn
 

braindecode/models/sleep_stager_eldele_2021.py

@@ -5,7 +5,6 @@
 import math
 import warnings
 from copy import deepcopy
-from typing import Callable, Optional
 
 import torch
 import torch.nn.functional as F

braindecode/models/sparcnet.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from collections import OrderedDict
 from math import floor, log2
-from typing import Any
 
 import torch
 import torch.nn as nn
@@ -69,6 +68,19 @@ class SPARCNet(EEGModuleMixin, nn.Module):
         conv_bias: bool = True,
         batch_norm: bool = True,
         activation: nn.Module = nn.ELU,
+        kernel_size_conv0: int = 7,
+        kernel_size_conv1: int = 1,
+        kernel_size_conv2: int = 3,
+        kernel_size_pool: int = 3,
+        stride_pool: int = 2,
+        stride_conv0: int = 2,
+        stride_conv1: int = 1,
+        stride_conv2: int = 1,
+        padding_pool: int = 1,
+        padding_conv0: int = 3,
+        padding_conv2: int = 1,
+        kernel_size_trans: int = 2,
+        stride_trans: int = 2,
         # EEGModuleMixin parameters
         # (another way to present the same parameters)
         chs_info=None,
@@ -96,9 +108,9 @@ class SPARCNet(EEGModuleMixin, nn.Module):
                     nn.Conv1d(
                         in_channels=self.n_chans,
                         out_channels=out_channels,
-                        kernel_size=7,
-                        stride=2,
-                        padding=3,
+                        kernel_size=kernel_size_conv0,
+                        stride=stride_conv0,
+                        padding=padding_conv0,
                         bias=conv_bias,
                     ),
                 )
@@ -106,7 +118,11 @@ class SPARCNet(EEGModuleMixin, nn.Module):
         )
         first_conv["norm0"] = nn.BatchNorm1d(out_channels)
         first_conv["act_layer"] = activation()
-        first_conv["pool0"] = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+        first_conv["pool0"] = nn.MaxPool1d(
+            kernel_size=kernel_size_pool,
+            stride=stride_pool,
+            padding=padding_pool,
+        )
 
         self.encoder = nn.Sequential(first_conv)
 
@@ -123,6 +139,11 @@ class SPARCNet(EEGModuleMixin, nn.Module):
                 conv_bias=conv_bias,
                 batch_norm=batch_norm,
                 activation=activation,
+                kernel_size_conv1=kernel_size_conv1,
+                kernel_size_conv2=kernel_size_conv2,
+                stride_conv1=stride_conv1,
+                stride_conv2=stride_conv2,
+                padding_conv2=padding_conv2,
             )
             self.encoder.add_module("denseblock%d" % (n_layer + 1), block)
             # update the number of channels after each dense block
@@ -134,16 +155,19 @@ class SPARCNet(EEGModuleMixin, nn.Module):
                 conv_bias=conv_bias,
                 batch_norm=batch_norm,
                 activation=activation,
+                kernel_size_trans=kernel_size_trans,
+                stride_trans=stride_trans,
             )
             self.encoder.add_module("transition%d" % (n_layer + 1), trans)
             # update the number of channels after each transition layer
             n_channels = n_channels // 2
 
+        self.adaptative_pool = nn.AdaptiveAvgPool1d(1)
+        self.activation_layer = activation()
+        self.flatten_layer = nn.Flatten()
+
         # add final convolutional layer
-        self.final_layer = nn.Sequential(
-            activation(),
-            nn.Linear(n_channels, self.n_outputs),
-        )
+        self.final_layer = nn.Linear(n_channels, self.n_outputs)
 
         self._init_weights()
 
@@ -178,7 +202,10 @@ class SPARCNet(EEGModuleMixin, nn.Module):
         torch.Tensor
             The output tensor of the model with shape (batch_size, n_outputs)
         """
-        emb = self.encoder(X).squeeze(-1)
+        emb = self.encoder(X)
+        emb = self.adaptative_pool(emb)
+        emb = self.activation_layer(emb)
+        emb = self.flatten_layer(emb)
         out = self.final_layer(emb)
         return out
 
@@ -224,6 +251,11 @@ class _DenseLayer(nn.Sequential):
         conv_bias: bool = True,
         batch_norm: bool = True,
         activation: nn.Module = nn.ELU,
+        kernel_size_conv1: int = 1,
+        kernel_size_conv2: int = 3,
+        stride_conv1: int = 1,
+        stride_conv2: int = 1,
+        padding_conv2: int = 1,
     ):
         super().__init__()
         if batch_norm:
@@ -235,8 +267,8 @@ class _DenseLayer(nn.Sequential):
             nn.Conv1d(
                 in_channels=in_channels,
                 out_channels=bottleneck_size * growth_rate,
-                kernel_size=1,
-                stride=1,
+                kernel_size=kernel_size_conv1,
+                stride=stride_conv1,
                 bias=conv_bias,
             ),
         )
@@ -248,9 +280,9 @@ class _DenseLayer(nn.Sequential):
             nn.Conv1d(
                 in_channels=bottleneck_size * growth_rate,
                 out_channels=growth_rate,
-                kernel_size=3,
-                stride=1,
-                padding=1,
+                kernel_size=kernel_size_conv2,
+                stride=stride_conv2,
+                padding=padding_conv2,
                 bias=conv_bias,
             ),
         )
@@ -311,6 +343,11 @@ class _DenseBlock(nn.Sequential):
         conv_bias=True,
         batch_norm=True,
         activation: nn.Module = nn.ELU,
+        kernel_size_conv1: int = 1,
+        kernel_size_conv2: int = 3,
+        stride_conv1: int = 1,
+        stride_conv2: int = 1,
+        padding_conv2: int = 1,
     ):
         super(_DenseBlock, self).__init__()
         for idx_layer in range(num_layers):
@@ -322,6 +359,11 @@ class _DenseBlock(nn.Sequential):
                 conv_bias=conv_bias,
                 batch_norm=batch_norm,
                 activation=activation,
+                kernel_size_conv1=kernel_size_conv1,
+                kernel_size_conv2=kernel_size_conv2,
+                stride_conv1=stride_conv1,
+                stride_conv2=stride_conv2,
+                padding_conv2=padding_conv2,
             )
             self.add_module(f"denselayer{idx_layer + 1}", layer)
 
@@ -360,6 +402,8 @@ class _TransitionLayer(nn.Sequential):
         conv_bias=True,
         batch_norm=True,
         activation: nn.Module = nn.ELU,
+        kernel_size_trans: int = 2,
+        stride_trans: int = 2,
     ):
         super(_TransitionLayer, self).__init__()
         if batch_norm:
@@ -375,4 +419,6 @@ class _TransitionLayer(nn.Sequential):
                 bias=conv_bias,
             ),
         )
-        self.add_module("pool", nn.AvgPool1d(kernel_size=2, stride=2))
+        self.add_module(
+            "pool", nn.AvgPool1d(kernel_size=kernel_size_trans, stride=stride_trans)
+        )

braindecode/models/tcn.py

@@ -8,7 +8,7 @@ from torch.nn import init
 from torch.nn.utils.parametrizations import weight_norm
 
 from braindecode.models.base import EEGModuleMixin
-from braindecode.modules import Chomp1d, Ensure4d, Expression, SqueezeFinalOutput
+from braindecode.modules import Chomp1d, Ensure4d, SqueezeFinalOutput
 
 
 class BDTCN(EEGModuleMixin, nn.Module):

braindecode/models/tsinception.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import torch
 import torch.nn as nn
 from einops.layers.torch import Rearrange
+from mne.utils import warn
 
 from braindecode.models.base import EEGModuleMixin
 
@@ -98,20 +99,44 @@ class TSceptionV1(EEGModuleMixin, nn.Module):
 
         ### Layers
         self.ensuredim = Rearrange("batch nchans time -> batch 1 nchans time")
+        if self.input_window_seconds < max(self.inception_windows):
+            inception_windows = (
+                self.input_window_seconds,
+                self.input_window_seconds / 2,
+                self.input_window_seconds / 4,
+            )
+            warning_msg = (
+                "Input window size is smaller than the maximum inception window size. "
+                "We are adjusting the input window size to match the maximum inception window size.\n"
+                f"Original input window size: {self.inception_windows}, \n"
+                f"Adjusted inception windows: {inception_windows}"
+            )
+            warn(warning_msg, UserWarning)
+            self.inception_windows = inception_windows
         # Define temporal convolutional layers (Tception)
-        self.temporal_blocks = nn.ModuleList(
-            [
-                self._conv_block(
-                    in_channels=1,
-                    out_channels=number_filter_temp,
-                    kernel_size=(1, int(window * self.sfreq)),
-                    stride=1,
-                    pool_size=self.pool_size,
-                    activation=self.activation,
-                )
-                for window in self.inception_windows
-            ]
-        )
+        self.temporal_blocks = nn.ModuleList()
+        for window in self.inception_windows:
+            # 1. Calculate the temporal kernel size for this block
+            kernel_size_t = int(window * self.sfreq)
+
+            # 2. Calculate the output length of the convolution
+            conv_out_len = self.n_times - kernel_size_t + 1
+
+            # 3. Ensure the pooling size is not larger than the conv output
+            #    and is at least 1.
+            dynamic_pool_size = max(1, min(self.pool_size, conv_out_len))
+
+            # 4. Create the block with the dynamic pooling size
+            block = self._conv_block(
+                in_channels=1,
+                out_channels=self.number_filter_temp,
+                kernel_size=(1, kernel_size_t),
+                stride=1,
+                pool_size=dynamic_pool_size,  # Use the dynamic size
+                activation=self.activation,
+            )
+            self.temporal_blocks.append(block)
+
         self.batch_temporal_lay = nn.BatchNorm2d(self.number_filter_temp)
 
         # Define spatial convolutional layers (Sception)

braindecode/models/util.py

@@ -5,11 +5,7 @@
 import inspect
 from pathlib import Path
 
-import numpy as np
 import pandas as pd
-import torch
-from scipy.special import log_softmax
-from sklearn.utils import deprecated
 
 import braindecode.models as models
 
@@ -76,12 +72,12 @@ models_mandatory_parameters = [
     ),  # 1 channel
     ("TIDNet", ["n_chans", "n_outputs", "n_times"], None),
     ("USleep", ["n_chans", "n_outputs", "n_times", "sfreq"], dict(sfreq=128.0)),
-    ("BIOT", ["n_chans", "n_outputs", "sfreq"], None),
+    ("BIOT", ["n_chans", "n_outputs", "sfreq", "n_times"], None),
     ("AttentionBaseNet", ["n_chans", "n_outputs", "n_times"], None),
     ("Labram", ["n_chans", "n_outputs", "n_times"], None),
     ("EEGSimpleConv", ["n_chans", "n_outputs", "sfreq"], None),
     ("SPARCNet", ["n_chans", "n_outputs", "n_times"], None),
-    ("ContraWR", ["n_chans", "n_outputs", "sfreq"], dict(sfreq=200.0)),
+    ("ContraWR", ["n_chans", "n_outputs", "sfreq", "n_times"], dict(sfreq=200.0)),
     ("EEGNeX", ["n_chans", "n_outputs", "n_times"], None),
     ("TSceptionV1", ["n_chans", "n_outputs", "n_times", "sfreq"], dict(sfreq=200.0)),
     ("EEGTCNet", ["n_chans", "n_outputs", "n_times"], None),

braindecode/modules/__init__.py

@@ -36,3 +36,49 @@ from .stats import (
 )
 from .util import aggregate_probas
 from .wrapper import Expression, IntermediateOutputWrapper
+
+__all__ = [
+    "LogActivation",
+    "SafeLog",
+    "CAT",
+    "CBAM",
+    "ECA",
+    "FCA",
+    "GCT",
+    "SRM",
+    "CATLite",
+    "EncNet",
+    "GatherExcite",
+    "GSoP",
+    "MultiHeadAttention",
+    "SqueezeAndExcitation",
+    "MLP",
+    "FeedForwardBlock",
+    "InceptionBlock",
+    "AvgPool2dWithConv",
+    "CausalConv1d",
+    "CombinedConv",
+    "Conv2dWithConstraint",
+    "DepthwiseConv2d",
+    "FilterBankLayer",
+    "GeneralizedGaussianFilter",
+    "Chomp1d",
+    "DropPath",
+    "Ensure4d",
+    "SqueezeFinalOutput",
+    "TimeDistributed",
+    "LinearWithConstraint",
+    "MaxNormLinear",
+    "MaxNorm",
+    "MaxNormParametrize",
+    "LogPowerLayer",
+    "LogVarLayer",
+    "MaxLayer",
+    "MeanLayer",
+    "StatLayer",
+    "StdLayer",
+    "VarLayer",
+    "aggregate_probas",
+    "Expression",
+    "IntermediateOutputWrapper",
+]

braindecode/modules/filter.py

@@ -1,18 +1,14 @@
 from __future__ import annotations
 
-from functools import partial
 from typing import Optional
 
 import torch
-from einops.layers.torch import Rearrange
 from mne.filter import _check_coefficients, create_filter
 from mne.utils import warn
 from torch import Tensor, from_numpy, nn
 from torch.fft import fftfreq
 from torchaudio.functional import fftconvolve, filtfilt
 
-import braindecode.functional as F
-
 
 class FilterBankLayer(nn.Module):
     """Apply multiple band-pass filters to generate multiview signal representation.