braindecode 1.5.0.dev1010__py3-none-any.whl → 1.5.0.dev1015__py3-none-any.whl

braindecode/augmentation/__init__.py

@@ -4,6 +4,7 @@ from . import functional
 from .base import AugmentedDataLoader, Compose, IdentityTransform, Transform
 from .transforms import (
     AmplitudeScale,
+    BandRotation,
     BandstopFilter,
     ChannelsDropout,
     ChannelsReref,
@@ -47,6 +48,7 @@ __all__ = [
     "SegmentationReconstruction",
     "MaskEncoding",
     "AmplitudeScale",
+    "BandRotation",
     "ChannelsReref",
     "functional",
 ]

braindecode/augmentation/functional.py

@@ -1298,3 +1298,130 @@ def amplitude_scale(
     X = s * X
 
     return X, y
+
+
+def band_rotation(
+    X: torch.Tensor,
+    y: torch.Tensor,
+    num_bands: int = 2,
+    electrodes_per_band: int = 16,
+    band_offsets: tuple[int, ...] = (-1, 0, 1),
+    max_temporal_jitter: int = 0,
+    circular_jitter: bool = True,
+    random_state: int | np.random.RandomState | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-band electrode rotation + inter-band temporal jitter.
+
+    Models small wristband rotation between sessions and relative timing
+    noise between two arms.  Introduced in [Sivakumar2024]_ for the
+    emg2qwerty CTC keystroke decoding task: each electrode band gets its
+    own circular roll along the channel axis (``Uniform(band_offsets)``
+    positions), and band 1 also gets a sample-level temporal shift
+    (``Uniform(-max_temporal_jitter, +max_temporal_jitter)``) along the
+    time axis.
+
+    Channel layout assumes ``(B, num_bands * electrodes_per_band, T)`` with
+    bands contiguous along the channel axis.  Same offset / shift is
+    applied to every sample in the batch (one set of parameters per call).
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        EMG input batch of shape ``(B, C, T)`` with
+        ``C == num_bands * electrodes_per_band``.
+    y : torch.Tensor
+        Labels (returned unchanged).
+    num_bands : int, optional
+        Number of electrode bands (e.g. ``2`` for left + right wristband).
+        Must be ``>= 1``.  Defaults to 2.
+    electrodes_per_band : int, optional
+        Electrodes per band (e.g. ``16``).  Must be ``>= 1``.  Defaults
+        to 16.
+    band_offsets : tuple of int, optional
+        Per-band roll values to sample from uniformly.  ``(-1, 0, 1)``
+        covers ±1-electrode misalignment.  Must be non-empty.  Defaults
+        to ``(-1, 0, 1)``.
+    max_temporal_jitter : int, optional
+        Max ±-sample temporal shift applied to band 1 only when
+        ``num_bands >= 2``.  Defaults to 0 (disabled).  Must be ``>= 0``.
+    circular_jitter : bool, optional
+        If True (the default, paper-faithful), the temporal jitter is a
+        circular ``torch.roll`` — samples shifted off one edge wrap to
+        the other.  If False, the gap left by the shift is zero-padded
+        and the shifted-off samples are dropped, avoiding wrap-around
+        discontinuity at the cost of a small zeroed margin.  Has no
+        effect when ``max_temporal_jitter == 0``.
+    random_state : int | numpy.random.RandomState, optional
+        Seed / generator for sampling rotation + jitter values.
+
+    Returns
+    -------
+    torch.Tensor
+        Transformed inputs.
+    torch.Tensor
+        Labels (unchanged).
+
+    References
+    ----------
+    .. [Sivakumar2024] Sivakumar, V., Seely, J., Du, A., Bittner, S. R.,
+       Berenzweig, A., Bolarinwa, A., Gramfort, A., & Mandel, M. I. (2024).
+       "emg2qwerty: A Large Dataset with Baselines for Touch Typing using
+       Surface Electromyography." *NeurIPS Datasets and Benchmarks Track*.
+    """
+    if num_bands < 1:
+        raise ValueError(f"num_bands must be >= 1, got {num_bands}")
+    if electrodes_per_band < 1:
+        raise ValueError(f"electrodes_per_band must be >= 1, got {electrodes_per_band}")
+    # Normalise to a tuple before truth-testing so callers can pass any
+    # sequence-like (incl. ``np.ndarray``) without hitting numpy's
+    # ambiguous-truth-value error on ``if not band_offsets``.
+    band_offsets = tuple(band_offsets)
+    if not band_offsets:
+        raise ValueError("band_offsets must be non-empty")
+    if not all(isinstance(o, (int, np.integer)) for o in band_offsets):
+        raise ValueError(f"band_offsets must contain integers, got {band_offsets!r}")
+    if max_temporal_jitter < 0:
+        raise ValueError(f"max_temporal_jitter must be >= 0, got {max_temporal_jitter}")
+    expected_channels = num_bands * electrodes_per_band
+    if X.shape[1] != expected_channels:
+        raise ValueError(
+            f"X.shape[1]={X.shape[1]} != num_bands * electrodes_per_band="
+            f"{expected_channels}"
+        )
+
+    rng = check_random_state(random_state)
+    band_offsets_arr = np.asarray(band_offsets)
+    out = X.clone()
+
+    # Per-band channel-axis rolls.  A vectorized ``torch.gather`` was
+    # benchmarked and is ~16 % slower for the typical ``num_bands == 2``
+    # case on CPU (the index tensor is larger than what two contiguous
+    # rolls touch); the gather only wins past ``num_bands >= 8``.
+    for b in range(num_bands):
+        offset = int(rng.choice(band_offsets_arr))
+        if offset:
+            sl = slice(b * electrodes_per_band, (b + 1) * electrodes_per_band)
+            out[:, sl, :] = torch.roll(out[:, sl, :], offset, dims=1)
+
+    # Inter-band temporal jitter — paper recipe applies it to band 1 only.
+    if max_temporal_jitter > 0 and num_bands >= 2:
+        shift = int(rng.randint(-max_temporal_jitter, max_temporal_jitter + 1))
+        if shift:
+            sl = slice(electrodes_per_band, 2 * electrodes_per_band)
+            band1 = out[:, sl, :]
+            if circular_jitter:
+                # Paper-faithful circular shift; wraps end-of-window
+                # samples to the start (and vice versa).
+                out[:, sl, :] = torch.roll(band1, shift, dims=2)
+            else:
+                # Crop-and-pad shift: drop samples that fall off one end,
+                # zero-pad the gap on the other.  Avoids the wrap-around
+                # discontinuity at the cost of a ``|shift|``-sample margin.
+                shifted = torch.zeros_like(band1)
+                if shift > 0:
+                    shifted[:, :, shift:] = band1[:, :, :-shift]
+                else:  # shift < 0
+                    shifted[:, :, :shift] = band1[:, :, -shift:]
+                out[:, sl, :] = shifted
+
+    return out, y

braindecode/augmentation/transforms.py

@@ -16,6 +16,7 @@ from mne.channels import make_standard_montage
 from .base import Transform
 from .functional import (
     amplitude_scale,
+    band_rotation,
     bandstop_filter,
     channels_dropout,
     channels_permute,
@@ -1356,3 +1357,99 @@ class AmplitudeScale(Transform):
     def get_augmentation_params(self, *batch):
         """Return transform parameters."""
         return {"random_state": self.rng, "scale": self.scale}
+
+
+class BandRotation(Transform):
+    """Per-band electrode rotation + inter-band temporal jitter.
+
+    Models small wristband rotation between sessions and relative timing
+    noise between two arms.  Introduced in [Sivakumar2024]_ for the
+    emg2qwerty surface-EMG keystroke decoding task: the channel axis is
+    laid out as ``(B, num_bands * electrodes_per_band, T)`` with bands
+    contiguous, each band gets a uniform circular roll along the channel
+    axis, and when ``num_bands >= 2``, band 1 also gets a sample-level
+    temporal shift.  The same offset / shift is applied to every sample
+    in a transformed sub-batch (one set of parameters per call).
+
+    Parameters
+    ----------
+    probability : float
+        Float setting the probability of applying the operation.
+    num_bands : int, optional
+        Number of electrode bands (e.g. ``2`` for left + right wristband).
+        Must be ``>= 1``.  Defaults to 2.
+    electrodes_per_band : int, optional
+        Electrodes per band (e.g. ``16``).  Must be ``>= 1``.  Defaults
+        to 16.
+    band_offsets : tuple of int, optional
+        Per-band roll values to sample from uniformly.  ``(-1, 0, 1)``
+        covers ±1-electrode misalignment.  Must be non-empty.  Defaults
+        to ``(-1, 0, 1)``.
+    max_temporal_jitter : int, optional
+        Max ±-sample temporal shift applied to band 1.  Defaults to 0
+        (jitter disabled).  Must be ``>= 0``.  The emg2qwerty paper uses
+        120 samples (60 ms at 2 kHz).
+    circular_jitter : bool, optional
+        If True (default, paper-faithful) the jitter is a circular roll;
+        if False the gap left by the shift is zero-padded.  See
+        :func:`band_rotation`.
+    random_state : int | numpy.random.RandomState, optional
+        Seed for the rotation / jitter sampler.  Defaults to None.
+
+    References
+    ----------
+    .. [Sivakumar2024] Sivakumar, V., Seely, J., Du, A., Bittner, S. R.,
+       Berenzweig, A., Bolarinwa, A., Gramfort, A., & Mandel, M. I. (2024).
+       "emg2qwerty: A Large Dataset with Baselines for Touch Typing using
+       Surface Electromyography." *NeurIPS Datasets and Benchmarks Track*.
+    """
+
+    operation = staticmethod(band_rotation)  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        probability,
+        num_bands=2,
+        electrodes_per_band=16,
+        band_offsets=(-1, 0, 1),
+        max_temporal_jitter=0,
+        circular_jitter=True,
+        random_state=None,
+    ):
+        super().__init__(probability=probability, random_state=random_state)
+        # Up-front parameter validation; the underlying ``band_rotation``
+        # also re-checks at call time, but raising here surfaces config
+        # mistakes when the Transform is built rather than on the first
+        # batch.
+        if num_bands < 1:
+            raise ValueError(f"num_bands must be >= 1, got {num_bands}")
+        if electrodes_per_band < 1:
+            raise ValueError(
+                f"electrodes_per_band must be >= 1, got {electrodes_per_band}"
+            )
+        band_offsets = tuple(band_offsets)
+        if not band_offsets:
+            raise ValueError("band_offsets must be non-empty")
+        if not all(isinstance(o, (int, np.integer)) for o in band_offsets):
+            raise ValueError(
+                f"band_offsets must contain integers, got {band_offsets!r}"
+            )
+        if max_temporal_jitter < 0:
+            raise ValueError(
+                f"max_temporal_jitter must be >= 0, got {max_temporal_jitter}"
+            )
+        self.num_bands = num_bands
+        self.electrodes_per_band = electrodes_per_band
+        self.band_offsets = band_offsets
+        self.max_temporal_jitter = max_temporal_jitter
+        self.circular_jitter = circular_jitter
+
+    def get_augmentation_params(self, *batch):
+        return {
+            "num_bands": self.num_bands,
+            "electrodes_per_band": self.electrodes_per_band,
+            "band_offsets": self.band_offsets,
+            "max_temporal_jitter": self.max_temporal_jitter,
+            "circular_jitter": self.circular_jitter,
+            "random_state": self.rng,
+        }

braindecode/models/emg2qwerty.py

@@ -56,7 +56,11 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
     Returns ``(batch, T_out, n_outputs)``. With ``n_times=8000`` and
     defaults, ``T_out=373``. For :class:`~torch.nn.CTCLoss`, transpose
     to ``(T_out, batch, n_outputs)``; use :meth:`compute_output_lengths`
-    for emission lengths.
+    for emission lengths. Pass ``return_features=True`` to return the
+    pre-classifier encoder representation as a
+    ``{"features": (batch, T_out, num_features), "cls_token": None}``
+    dict, matching the BIOT / signal-JEPA convention used by downstream
+    wrappers (e.g. neuroai's ``DownstreamWrapperModel``).
 
     .. rubric:: Paper training recipe
 
@@ -69,7 +73,9 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
       local minimum).
     - **Augmentation**: per-band electrode rotations by -1/0/+1 positions,
       ±60-sample temporal jitter, and SpecAugment [park2019specaug]_ on
-      the log-spectrogram.
+      the log-spectrogram. SpecAugment is built into the model
+      (``spec_augment=True``) and only fires in training mode; the
+      time/frequency-jitter pieces are dataset-side augmentations.
     - **Decoding**: greedy CTC. Upstream also reports a 6-gram KenLM
       beam decoder, not ported here.
 
@@ -145,6 +151,37 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
         layers and again after the second :class:`~torch.nn.Linear`.
         Default ``0.0`` matches the upstream paper recipe (no dropout).
         Set ``> 0`` for regularized training.
+    spec_augment : bool
+        If ``True``, apply SpecAugment [park2019specaug]_ time/frequency
+        masking on the log-spectrogram during training only. Disabled in
+        ``eval`` mode and absent from the parameter / state-dict count.
+        Defaults to ``False``; set to ``True`` to match the upstream
+        emg2qwerty paper recipe.
+    n_time_masks : int
+        Maximum number of time masks applied per call. Each forward pass
+        samples a uniform integer in ``[0, n_time_masks]``. Defaults to
+        ``3`` (Sivakumar et al. Sec 5.2).
+    time_mask_param : int
+        Maximum time-mask width in spectrogram frames. Defaults to ``25``.
+    n_freq_masks : int
+        Maximum number of frequency masks applied per call. Each forward
+        pass samples a uniform integer in ``[0, n_freq_masks]``. Defaults
+        to ``2``.
+    freq_mask_param : int
+        Maximum frequency-mask width in STFT bins. Defaults to ``4``.
+    spec_augment_prob : float
+        Probability of running SpecAugment on a given training batch
+        (Bernoulli gate before sampling mask counts). Defaults to ``1.0``.
+    return_feature : bool
+        If ``True``, ``forward`` returns a tuple
+        ``(emissions, features)`` instead of just ``emissions`` —
+        :class:`braindecode.models.BIOT`-style legacy feature path. Lets
+        configuration-driven downstream wrappers (e.g. neuroai's
+        ``DownstreamWrapperModel`` with ``model_output_key=1``) pick up
+        the encoder representation without passing a runtime kwarg.
+        Defaults to ``False``. Mutually compatible with the runtime
+        ``return_features`` (plural) flag, which still wins when set
+        to ``True``.
 
     Examples
     --------
@@ -216,6 +253,13 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
         log_softmax: bool = False,
         activation: type[nn.Module] = nn.ReLU,
         drop_prob: float = 0.0,
+        spec_augment: bool = False,
+        n_time_masks: int = 3,
+        time_mask_param: int = 25,
+        n_freq_masks: int = 2,
+        freq_mask_param: int = 4,
+        spec_augment_prob: float = 1.0,
+        return_feature: bool = False,
         # Standard braindecode args
         n_times: int | None = None,
         input_window_seconds: float | None = None,
@@ -256,6 +300,7 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
         self.hop_length = hop_length
         self.kernel_width = kernel_width
         self.log_softmax = log_softmax
+        self.return_feature = return_feature
 
         n_freq_bins = n_fft // 2 + 1
         in_features = electrodes_per_band * n_freq_bins
@@ -269,6 +314,23 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
             log_eps=log_eps,
         )
 
+        # Built-in SpecAugment lives between the spectrogram and the BatchNorm
+        # so it operates on the log-power tensor (matches upstream emg2qwerty
+        # and the previous neuralbench callback). ``nn.Identity`` keeps the
+        # forward path symmetrical without contributing parameters or
+        # state-dict keys when SpecAugment is disabled.
+        self.spec_augment: nn.Module
+        if spec_augment:
+            self.spec_augment = _SpecAugment(
+                n_time_masks=n_time_masks,
+                time_mask_param=time_mask_param,
+                n_freq_masks=n_freq_masks,
+                freq_mask_param=freq_mask_param,
+                prob=spec_augment_prob,
+            )
+        else:
+            self.spec_augment = nn.Identity()
+
         # Indices 0/1/3 match upstream's ``TDSConvCTCModule.model``;
         # index 2 is a parameter-free Flatten; upstream's index 4 (head)
         # is broken out as ``self.final_layer`` and remapped via :attr:`mapping`.
@@ -298,7 +360,13 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
             isinstance(m, _TDSConv2dBlock) for m in self.model[3].tds_conv_blocks
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: torch.Tensor, return_features: bool = False
+    ) -> (
+        torch.Tensor
+        | dict[str, torch.Tensor | None]
+        | tuple[torch.Tensor, torch.Tensor]
+    ):
         """Run the full pipeline.
 
         Parameters
@@ -307,12 +375,37 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
             Raw EMG of shape ``(batch, n_chans=32, n_times)``. ``n_times``
             must be at least the encoder's receptive field, ``n_fft +
             n_conv_blocks * (kernel_width - 1) * hop_length``.
+        return_features : bool
+            If ``True``, return a ``dict`` with the encoder representation
+            instead of the classification emissions. The encoder is the
+            full TDS-Conv stack up to (but not including)
+            ``self.final_layer`` — i.e. what downstream wrappers want
+            when they apply their own probe/aggregation. Matches the
+            BIOT / signal-JEPA convention so the same neuroai
+            ``DownstreamWrapperModel(model_output_key="features")``
+            can consume it. Wins over the constructor-time
+            ``return_feature`` flag when set.
 
         Returns
         -------
-        emissions : torch.Tensor
-            Shape ``(batch, T_out, n_outputs)``. Log-probabilities if
+        torch.Tensor or dict or tuple
+            Default (``return_features=False``, init
+            ``return_feature=False``): ``torch.Tensor`` of shape
+            ``(batch, T_out, n_outputs)``. Log-probabilities if
             ``log_softmax=True``, otherwise logits.
+
+            If runtime ``return_features=True``: ``dict`` with
+            ``"features"`` (shape ``(batch, T_out, num_features)``,
+            where ``num_features = num_bands * mlp_features[-1]``) and
+            ``"cls_token"`` (always ``None`` — TDS-Conv has no
+            ``[CLS]``).
+
+            If init ``return_feature=True`` and runtime
+            ``return_features=False``: tuple ``(emissions, features)``
+            where ``features`` has shape ``(batch, T_out,
+            num_features)``. Same layout BIOT exposes for
+            configuration-driven feature extraction (e.g. neuroai's
+            ``model_output_key=1``).
         """
         if x.ndim != 3 or x.shape[-2] != self.n_chans:
             raise ValueError(
@@ -331,11 +424,24 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
                 f"kernel_width={self.kernel_width})."
             )
         spectrogram = self.spectrogram(x)
+        spectrogram = self.spec_augment(spectrogram)
         encoded = self.model(spectrogram)
+        # ``encoded`` is (T_out, B, num_features); only materialise the
+        # batch-first features tensor in the branches that actually return
+        # it, so the default emissions-only path skips the extra transpose
+        # + contiguous copy on every forward.
+        if return_features:
+            return {
+                "features": encoded.transpose(0, 1).contiguous(),
+                "cls_token": None,
+            }
         emissions = self.final_layer(encoded)
         if self.log_softmax:
             emissions = F.log_softmax(emissions, dim=-1)
-        return emissions.transpose(0, 1).contiguous()
+        emissions = emissions.transpose(0, 1).contiguous()
+        if self.return_feature:
+            return emissions, encoded.transpose(0, 1).contiguous()
+        return emissions
 
     def reset_head(self, n_outputs: int) -> None:
         """Replace the classification head for a new vocabulary size.
@@ -411,7 +517,13 @@ class EMG2QwertyNet(EEGModuleMixin, nn.Module):
             dummy_input = torch.zeros(
                 1, self.n_chans, n_times, dtype=dtype, device=device
             )
-            return tuple(self.forward(dummy_input).shape)
+            # ``return_features=False`` keeps the dict path off; the init
+            # ``return_feature`` flag may still produce a tuple, so unpack
+            # the emissions explicitly to report the public output shape.
+            out = self.forward(dummy_input, return_features=False)
+            emissions = out[0] if isinstance(out, tuple) else out
+            assert isinstance(emissions, torch.Tensor)
+            return tuple(emissions.shape)
 
 
 class _LogSpectrogram(nn.Module):
@@ -483,6 +595,94 @@ class _LogSpectrogram(nn.Module):
         ).movedim(-1, 0)
 
 
+class _SpecAugment(nn.Module):
+    r"""SpecAugment masking on the log-spectrogram during training.
+
+    Applies up to ``n_time_masks`` × ``time_mask_param``-frame time
+    bands and ``n_freq_masks`` × ``freq_mask_param``-bin frequency
+    bands. Masks are independent per ``(sample × band × electrode)``
+    triple — same recipe as the upstream emg2qwerty
+    :class:`emg2qwerty.transforms.SpecAugment` dataset transform
+    (Sivakumar et al. Sec 5.2 / NeurIPS 2024), which is
+    :func:`torchaudio.functional.mask_along_axis_iid`-style masking
+    sampled per leading dim of a spectrogram with shape
+    ``(..., freq, time)``. No-op outside ``training``.
+
+    The mask fill value is the on-device mean of the spectrogram —
+    ``log(power=1)=0`` would sit well above the typical log-power
+    distribution and inject artificial spikes — and stays a 0-D
+    tensor so the forward pass adds no host round-trip on GPU.
+    """
+
+    def __init__(
+        self,
+        n_time_masks: int = 3,
+        time_mask_param: int = 25,
+        n_freq_masks: int = 2,
+        freq_mask_param: int = 4,
+        prob: float = 1.0,
+    ) -> None:
+        super().__init__()
+        if n_time_masks < 0 or n_freq_masks < 0:
+            raise ValueError(
+                f"n_time_masks and n_freq_masks must be >= 0; got "
+                f"n_time_masks={n_time_masks}, n_freq_masks={n_freq_masks}."
+            )
+        if time_mask_param < 0 or freq_mask_param < 0:
+            raise ValueError(
+                f"time_mask_param and freq_mask_param must be >= 0; got "
+                f"time_mask_param={time_mask_param}, "
+                f"freq_mask_param={freq_mask_param}."
+            )
+        if not 0.0 <= prob <= 1.0:
+            raise ValueError(f"prob must be in [0, 1]; got {prob}.")
+        self.n_time_masks = n_time_masks
+        self.time_mask_param = time_mask_param
+        self.n_freq_masks = n_freq_masks
+        self.freq_mask_param = freq_mask_param
+        self.prob = prob
+        # ``iid_masks=True`` so masking is sampled over every leading dim
+        # except the trailing ``(freq, time)`` pair — i.e. one mask per
+        # ``(sample × band × electrode)`` on a 5-D
+        # ``(B, num_bands, electrodes, freq, T)`` input. Matches upstream
+        # emg2qwerty's per-``(band × electrode)`` dataset-time recipe.
+        self.time_mask = ta_transforms.TimeMasking(time_mask_param, iid_masks=True)
+        self.freq_mask = ta_transforms.FrequencyMasking(freq_mask_param, iid_masks=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # ``x``: (T_spec, B, num_bands, electrodes, freq).
+        if (
+            not self.training
+            or self.prob <= 0.0
+            or (self.n_time_masks == 0 and self.n_freq_masks == 0)
+        ):
+            return x
+        # All RNG draws use ``x.device`` so reproducibility seeds the same
+        # stream regardless of whether the user calls ``torch.manual_seed``
+        # or ``torch.cuda.manual_seed`` — and so torchaudio's internal
+        # device-side RNG and our Python-level gate stay in sync. ``.item()``
+        # still forces a host sync for the Python ``if``/loop bound, but
+        # that is unavoidable for control flow.
+        if self.prob < 1.0 and torch.rand((), device=x.device).item() >= self.prob:
+            return x
+        # ``torchaudio`` masking expects ``(..., freq, time)``; here that means
+        # ``(B, num_bands, electrodes, freq, T_spec)``. Move time to the end
+        # rather than reshaping into 4D, because ``mask_along_axis_iid`` draws
+        # one mask per leading-axis index, so the 5-D layout already gives the
+        # desired per-``(B × num_bands × electrodes)`` independence.
+        spec = x.movedim(0, -1).contiguous()
+        # 0-D on-device tensor — ``masked_fill`` / ``torch.where`` accept it
+        # without a host sync.
+        mask_value = spec.mean()
+        n_t = int(torch.randint(self.n_time_masks + 1, (), device=x.device).item())
+        for _ in range(n_t):
+            spec = self.time_mask(spec, mask_value=mask_value)
+        n_f = int(torch.randint(self.n_freq_masks + 1, (), device=x.device).item())
+        for _ in range(n_f):
+            spec = self.freq_mask(spec, mask_value=mask_value)
+        return spec.movedim(-1, 0)
+
+
 class _SpectrogramNorm(nn.Module):
     r""":class:`~torch.nn.BatchNorm2d` over (band × electrode) channels.
 

braindecode/version.py

@@ -1 +1 @@
-__version__ = "1.5.0.dev1010"
+__version__ = "1.5.0.dev1015"

{braindecode-1.5.0.dev1010.dist-info → braindecode-1.5.0.dev1015.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev1010
+Version: 1.5.0.dev1015
 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>

{braindecode-1.5.0.dev1010.dist-info → braindecode-1.5.0.dev1015.dist-info}/RECORD

@@ -3,11 +3,11 @@ braindecode/classifier.py,sha256=7kC_oY_UzHEes_WWdCvEpiA1ZKxMeuLL5tIPp5rfcpg,962
 braindecode/eegneuralnet.py,sha256=xjE6aPZdCQPs29NIpy_m1GLMMC2WZ3Db0Fuh1-xE1h4,13827
 braindecode/regressor.py,sha256=KiMJpqCUPWA2k2JWk9HGYTzeoBqJ4gAKEudeUVcFZY4,9266
 braindecode/util.py,sha256=f8bNIwt-SwsHqheH_BADQxTtA9oPt3Lb7GFnoI-Huwc,14101
-braindecode/version.py,sha256=dhwXuCj7AR3Ot0DD1Q7xVlLGgTrkwdesqS9QyInsZcM,30
-braindecode/augmentation/__init__.py,sha256=4xune2QUK6KHMKsAqijF7I9eeiVbP0wEoQJjCNLNcKM,1081
+braindecode/version.py,sha256=H7kXWUs3T_eeGV04VfvDHjBTgxYEXSFgslEWx9TudIs,30
+braindecode/augmentation/__init__.py,sha256=hmnjUsL_DX5BxYVdyNReh7T3YRQEJKYzciB1UwYHRvc,1119
 braindecode/augmentation/base.py,sha256=OJ1shOljI1yTY9zh2qWxQwivlY43sfx9Q-MAyMhxtPs,7338
-braindecode/augmentation/functional.py,sha256=q2k6mAXrujYlOZUndcjZN8e8b-6oJF1gGsORAI23hyE,43998
-braindecode/augmentation/transforms.py,sha256=x-3pwX0PtMHfSnPLGKNXbpTSk7j17Ci2FG_-646scg4,47268
+braindecode/augmentation/functional.py,sha256=jGKTNgWf9ZIFGcShYD0Qlb9IuC47OIwC813eCEcTPsM,49757
+braindecode/augmentation/transforms.py,sha256=QPS-cjbHz0TcKbd5Uuiag0s2Kt83xDhH1juBvAK6C5M,51426
 braindecode/datasets/__init__.py,sha256=rVOBadwqYBiMz5kl7nGiBOmMgr11xvjS4nuzzZTOn1U,1102
 braindecode/datasets/base.py,sha256=3lKLZQO4hfA-dv_JJEfPwyZ5nzRkLTu4qiRAqFVZUUQ,70508
 braindecode/datasets/bbci.py,sha256=SCm7OnCObotILQ0B1EdmZPoyJtzsRXpeU_gNKtqQLSc,19288
@@ -66,7 +66,7 @@ braindecode/models/eegpt.py,sha256=5ZXItSURum1vfUciuCHbJGIkXLk6XISP-e3NsFIv3wc,4
 braindecode/models/eegsimpleconv.py,sha256=suHO-v9laImwvXpLF2dwvoFFBKjiV-czAW1FHwRSscI,7306
 braindecode/models/eegsym.py,sha256=-5wb28oxx3YSCkFUnla-6P0RdGYshBPhfke7vSj-tnA,34592
 braindecode/models/eegtcnet.py,sha256=awEIwEIWSvS0b2Hb7ROfxV9DSwNe5z2224a-Teznuyo,10916
-braindecode/models/emg2qwerty.py,sha256=n_6MYSZywr3W4Id26OazI3mJHDQc6qumQF0TGk6wTBQ,28897
+braindecode/models/emg2qwerty.py,sha256=ln5Gmf7u0dup4_PN6xLRXs0KY-TrRb606p0s6l5J0o8,38925
 braindecode/models/fbcnet.py,sha256=YE5pCtF0Oo3J7rh8DDBl0oYZy9Tb2oyXkOYJJMr76Bo,7711
 braindecode/models/fblightconvnet.py,sha256=bOo7DlFiqByVQ0e6ethv5n2J7N-tIhiasLObqGLAg4g,11107
 braindecode/models/fbmsnet.py,sha256=prw9LcZBH_mEwV__fhUOOTbK4bmRdoKLLpjNuLA94Yg,12355
@@ -128,9 +128,9 @@ braindecode/visualization/frequency.py,sha256=gNwkn9yIik5SUp7d9HE9J_vPVGyzNsxxCO
 braindecode/visualization/metrics.py,sha256=j01kc04P9uEkQ2g2Tt2C76yr6soIj31PAuBMflrmODg,13615
 braindecode/visualization/sanity.py,sha256=nNClauUC8dCj_KCy_1RmaPDQAqExLczfPtUeQ7k9-Q0,4812
 braindecode/visualization/topology.py,sha256=mXxUfCCUJqa_cMF4y6GC3_A-qBCcS4uTc0EzBolkytE,2274
-braindecode-1.5.0.dev1010.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
-braindecode-1.5.0.dev1010.dist-info/licenses/NOTICE.txt,sha256=ZFFhigxIaKgDcMjCzPyAVSFV42ztU0kLOENt_kvherw,857
-braindecode-1.5.0.dev1010.dist-info/METADATA,sha256=4sQBBGOi3h1EE-DCZxdlh3Hswxjd4jQJ_fTdFBUSsyc,10275
-braindecode-1.5.0.dev1010.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-braindecode-1.5.0.dev1010.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
-braindecode-1.5.0.dev1010.dist-info/RECORD,,
+braindecode-1.5.0.dev1015.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
+braindecode-1.5.0.dev1015.dist-info/licenses/NOTICE.txt,sha256=ZFFhigxIaKgDcMjCzPyAVSFV42ztU0kLOENt_kvherw,857
+braindecode-1.5.0.dev1015.dist-info/METADATA,sha256=6-m5pOAtFg3doXKj78dFw_JSObw3BgLBtZBy_0gYhYI,10275
+braindecode-1.5.0.dev1015.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+braindecode-1.5.0.dev1015.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
+braindecode-1.5.0.dev1015.dist-info/RECORD,,