Support higher-dimensional data in `mne.decoding.Scaler` and `mne.decoding.Vectorizer`

Describe the new feature or enhancement

mne.decoding.Scaler and mne.decoding.Vectorizer assume inputs are 3D (i.e. from an Epochs).

However, this makes it impossible to use them with data from an EpochsTFR, especially in a pipeline with mne.decoding.SlidingEstimator (which needs data to have a time dimension).

Describe your proposed implementation

If input arrays are 3D, preserve the current behavior.

If the input arrays are 4D (corresponding to EpochsTFR):

• in the Scaler, transpose+reshape the data into (n_channels * n_frequencies, n_epochs * n_times)
• in the Vectorizer, reshape the data into (n_epochs, n_channels, -1)

Here is working code that lets me do this (assuming inputs are valid NumPy arrays):

import typing
from typing import Self

from sklearn.base import BaseEstimator, MetaEstimatorMixin, TransformerMixin

if typing.TYPE_CHECKING:
    import numpy as np
    import numpy.typing as npt
    from sklearn.preprocessing import RobustScaler, StandardScaler


class ChannelScaler(BaseEstimator, TransformerMixin, MetaEstimatorMixin):
    def __init__(
        self,
        base_scaler: StandardScaler | RobustScaler,
    ) -> None:
        self.shape_: tuple[int]
        self.base_scaler = base_scaler

    def fit(
        self,
        X: npt.NDArray[np.floating],  # ruff: ignore[invalid-argument-name]
        y: None = None,  # ruff: ignore[unused-method-argument]
    ) -> Self:
        self.shape_ = X.shape
        self.base_scaler.fit(self._reshape_to_sklearn(X))
        return self

    def transform(
        self,
        X: npt.NDArray[np.floating],  # ruff: ignore[invalid-argument-name]
        y: None = None,  # ruff: ignore[unused-method-argument]
    ) -> npt.NDArray[np.floating]:
        x_transformed = self.base_scaler.transform(self._reshape_to_sklearn(X))
        return self._restore_from_sklearn(x_transformed)

    def fit_transform(
        self,
        X: npt.NDArray[np.floating],  # ruff: ignore[invalid-argument-name]
        y: None = None,  # ruff: ignore[unused-method-argument]
        **fit_kws,  # ruff: ignore[missing-type-kwargs, unused-method-argument]
    ) -> npt.NDArray[np.floating]:
        self.shape_ = X.shape
        x_reshaped = self._reshape_to_sklearn(X)
        self.base_scaler.fit(x_reshaped)
        return self._restore_from_sklearn(self.base_scaler.transform(x_reshaped))

    @staticmethod
    def _reshape_to_sklearn(
        x: npt.NDArray[np.floating],
    ) -> npt.NDArray[np.floating]:
        return (
            x
            # (epochs, ..., times) -> (epochs, times, ...)
            .transpose(0, -1, *range(1, x.ndim - 1))
            # (epochs, times, ...) -> (epochs * times, -1)
            .reshape(x.shape[0] * x.shape[-1], -1)
        )

    def _restore_from_sklearn(
        self,
        x: npt.NDArray[np.floating],
    ) -> npt.NDArray[np.floating]:
        n_epochs = int(x.shape[0] / self.shape_[-1])
        return (
            x
            # (epochs * times, -1) -> (epochs, times, ...)
            .reshape(n_epochs, self.shape_[-1], *self.shape_[1:-1])
            # (epochs, times, ...) -> (epochs, ..., times)
            .transpose(0, *range(2, len(self.shape_)), 1)
        )


class Vectorizer(BaseEstimator, TransformerMixin):
    """
    Transform (n_samples, n_features, ...) array into 3D array of (n_samples, n_features, -1).

    This class reshapes an n-dimensional array into an (n_samples, n_features,
    -1) array, usable by the estimators and transformers of scikit-learn, in
    conjunction with mne.decoding.SlidingEstimator and
    mne.decoding.GeneralizingEstimator, which support parallelization along the
    last dimension.
    """

    def fit(
        self,
        X: npt.NDArray[np.floating],  # ruff: ignore[invalid-argument-name]
        y: None = None,  # ruff: ignore[unused-method-argument]
    ) -> Self:
        self.features_shape_ = X.shape[2:]
        return self

    @typing.override
    def transform(self, X: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
        """Reshape (n_samples, n_features, ...) array to (n_samples, n_features, -1).

        Parameters
        ----------
        X : array, shape (n_samples, n_features, ...)
            The data to fit. Must be array of at least 3D.

        Returns
        -------
        X : array, shape (n_samples, n_features, -1)
            The transformed data.
        """
        return X.reshape(*X.shape[:2], -1)

    def inverse_transform(
        self,
        X: npt.NDArray[np.floating],  # ruff: ignore[invalid-argument-name]
    ) -> npt.NDArray[np.floating]:
        """Reshape (n_samples, n_features, -1) back to its original feature shape.

        Parameters
        ----------
        X : array, shape (n_samples, n_features, -1)
            Data to be transformed back to original shape.

        Returns
        -------
        X : array, shape (n_samples, n_features, ...)
            The data transformed into shape as used in fit.
        """
        return X.reshape(*X.shape[:2], *self.features_shape_)

Describe possible alternatives

1. Leave it alone: the user can just implement their own Scaler when working with higher-dimensional data, as I've done.
2. Create a new class (e.g. mne.decoding.TemporalScaler) with this behavior.

Additional context

Happy to try and submit a PR if you're willing!

Start with the mne.decoding.Scaler and mne.decoding.Vectorizer entry points and inspect how they currently handle 3D arrays. Add the proposed 4D EpochsTFR handling while preserving existing 3D behavior, then verify that the reshaped output works with mne.decoding.SlidingEstimator. Done means higher-dimensional data can pass through the pipeline with the intended channel, frequency, epoch, and time dimensions.