Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraMod, LUNA; ZUNA default to pad by bruAristimunha · Pull Request #1226 · braindecode/braindecode

bruAristimunha · 2026-10-05T10:43:49Z

Summary

Four pretrained models raised on a window that is not a multiple of patch_size (LaBraM: reshape error, CBraMod/LUNA: einops, ZUNA: on_non_divisible="error" default). They now pad the last patch with zeros by default (one warning), through the shared PatchTokenizer, and raise only with on_non_divisible="error".

First step of a short stack (next: a geometry-contract test over every pretrained model; then channel handling). Design discussion in the linked issue.

Changes

labram.py: _SegmentPatch holds a parameter-free PatchTokenizer (pads before the conv/reshape patcher); new on_non_divisible kwarg on Labram. Decoder mode (_PatchEmbed, pre-patched input) untouched.
cbramod.py: the Rearrange becomes a PatchTokenizer; _n_patch() sizes the flat head with padding.
luna.py: tokenizer in _PatchEmbedNetwork and prepare_tokens (mask padded alongside the signal).
zuna.py: default on_non_divisible "error" → "pad"; its test asserts the new default and the explicit opt-out.
test_foundation_models.py: one check per model — pads with a warning, forwards finite, adds no state-dict key, raises with "error".
docs/whats_new.rst.

Testing

pytest test/unit_tests/models/test_foundation_models.py test/unit_tests/models/test_integration.py -k "labram or cbramod or luna or zuna" → 119 passed, 17 skipped.
Parity vs base 2426ddc3 on divisible windows (seeded weights, same input): state-dict keys identical, outputs bit-identical (max |Δ| = 0.0) for all four → released checkpoints unaffected.
ruff format / ruff check (0.14.9) on changed files: clean (two F841 pre-existing on master untouched).
Not run: full suite (CI).
Regression tests
Style checks recorded
docs/whats_new.rst updated

Notes for reviewers

Behaviour change: pad instead of raise for these four models (ZUNA's default flips). on_non_divisible="error" restores strictness per model. PatchTokenizer has no parameters, so no checkpoint conversion.

Copilot

Copilot review overview

🟡 Changes recommended

The contract mishandles expected-error xfails, while the audit can report non-finite outputs and incorrect channel modalities as successful.

Review effort: Balanced
Findings: 2 · 1

Open (3)

Preserve declared channel modality in synthetic geometries · New
Add alternate sampling-rate cases to geometry contract · New
Update migration status to reflect implemented work · New

What changed in this PR

Adds default padding for non-divisible windows in pretrained foundation models and introduces compatibility auditing.

Changes:

Integrates PatchTokenizer into LaBraM, CBraMod, and LUNA; changes ZUNA’s default to padding.
Adds geometry contract tests and a Hub checkpoint audit script.
Documents the compatibility design and behavior changes.

File	Description
`braindecode/models/labram.py`	Supports padding, cropping, or rejecting partial patches.
`braindecode/models/cbramod.py`	Replaces reshaping with `PatchTokenizer`.
`braindecode/models/luna.py`	Tokenizes signals and masks consistently.
`braindecode/models/zuna.py`	Defaults non-divisible windows to padding.
`test/unit_tests/models/test_pretrained_compat.py`	Adds the geometry compatibility contract.
`test/unit_tests/models/test_foundation_models.py`	Updates ZUNA behavior tests and formatting.
`scripts/compat_audit.py`	Audits random and pretrained models across geometries.
`docs/design/pretrained_compatibility.md`	Defines the broader compatibility proposal.
`docs/whats_new.rst`	Records the new behavior and tooling.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+def chs_coords_only(n):
+    """n channels with synthetic head-frame coordinates and non-standard names."""
+    chs = []
+    for i in range(n):
+        th = 2 * math.pi * i / n
+        loc = np.zeros(12)
+        loc[:3] = [0.08 * math.cos(th), 0.08 * math.sin(th), 0.03]
+        chs.append({"ch_name": f"E{i + 1}", "kind": "eeg", "loc": loc})


+    if sf:
+        geos["G5a"] = dict(chs_info=base, sfreq=sf, n_times=int(sf))
+        geos["G5b"] = dict(chs_info=base, sfreq=sf, n_times=int(sf * 30))
+        geos["G5c"] = dict(chs_info=base, sfreq=sf, n_times=nt + 37)


+Status: proposal, branch `feat/pretrained-compat-d`, nothing merged. Companion audit:
+`scripts/compat_audit.py` (geometry grid × 19 pretrained classes, random weights).


bruAristimunha · 2026-10-05T11:18:43Z

Self-review (lazy-senior-dev pass) — the first push of this PR was three PRs in one coat plus two files that don't belong in the library. Split:

this PR → the model change only: PatchTokenizer in LaBraM/CBraMod/LUNA, ZUNA default pad, one check per model in test_foundation_models.py. 6 files, +198/−42. Reverted the LaBraM decoder-mode (_PatchEmbed) edit: its input is pre-patched, nothing to pad there — YAGNI.
Geometry-compatibility contract test for every model with released weights (stacked on #1226) #1228 (stacked) → the contract test over all 19 pretrained classes. Dropped its two helper tests (now covered here) and a dead length field.
Pretrained models compatible by default: input-geometry contract, audits of the 19 pretrained classes, stacked plan #1227 → the design doc and both audit matrices (class-level and from_pretrained-level). Removed docs/design/*.md and scripts/compat_audit.py from the repo: the test is the mechanism; a script that duplicates its grid is a second thing to keep in sync.

Still to come as their own small PRs (see #1227): the from_pretrained channel-count check that blocks EEGPT/STEEGFormer/Brant/MVPFormer on any chs_info, the CBraMod from_pretrained LazyLinear failure, ChannelTokenizer.

Parity after the split: outputs bit-identical to base on divisible windows, same state-dict keys, 119 passed on the four models' tests.

…od, LUNA; ZUNA default on_non_divisible='pad'

Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 4 · 4

Open (8)

Same concern as in LUNA: _prepare_input is a private method. To avoid relying on tokenizer… · New
_prepare_input is a private method (leading underscore). Calling it from outside PatchTokenizer… · New
Add alternate sampling-rate cases to geometry contract
Preserve declared channel modality in synthetic geometries
self.rearrange is now a PatchTokenizer, not an einops Rearrange, so the attribute name is… · New
The note says non-divisible windows 'previously ... raised a reshape error', but the PR description… · New
This assertion is a bit brittle because state_dict() includes buffers (not just parameters) and… · New
Update migration status to reflect implemented work

        X_patch: Tensor
            [batch, n_chans, n_times//patch_size, patch_size]
        """
+        x = self.tokenizer._prepare_input(x)


+        x_signal = self.tokenizer._prepare_input(x_signal)
+        if mask is not None:
+            mask = self.tokenizer._prepare_input(mask)


        del n_chans, chs_info, n_times, input_window_seconds, sfreq, n_outputs
-        self.rearrange = Rearrange(
-            "batch n_chans (n_patch patch_size) -> batch n_chans n_patch patch_size",
+        # Shared tokenizer: (batch, n_chans, n_times) -> (batch, n_chans, n_patch, patch_size),
+        # padding/cropping a non-divisible time axis at forward time.
+        self.rearrange = PatchTokenizer(
            patch_size=patch_size,
+            n_times=self._n_times if self._n_times is not None else patch_size,
+            on_non_divisible=on_non_divisible,
        )


+  :class:`braindecode.modules.PatchTokenizer` and accept a window that is not a
+  multiple of ``patch_size`` by right zero-padding the last patch (with a warning),
+  through a new ``on_non_divisible={"pad", "crop", "error"}`` argument; previously
+  such windows raised a reshape error. :class:`braindecode.models.ZUNA` now defaults to


+    with torch.no_grad():
+        out = model(x, ch_names=_TEN_TWENTY) if cls is Labram else model(x)
+    assert torch.isfinite(out).all()
+    assert not any("tokenizer" in k for k in model.state_dict())


Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 7 · 5

Open (12)

_SegmentPatch also calls the private PatchTokenizer._prepare_input. To avoid locking in… · New
The PR description says decoder mode (_PatchEmbed, pre-patched input) is untouched, but this diff… · New
_prepare_input is a private method (leading underscore) of PatchTokenizer, so depending on it… · New
_prepare_input is a private method (leading underscore). Calling it from outside PatchTokenizer…
Same concern as in LUNA: _prepare_input is a private method. To avoid relying on tokenizer…
Add alternate sampling-rate cases to geometry contract
Preserve declared channel modality in synthetic geometries
self.rearrange now holds a PatchTokenizer rather than an einops Rearrange layer. This is… · New
This assertion is a bit brittle because state_dict() includes buffers (not just parameters) and…
The note says non-divisible windows 'previously ... raised a reshape error', but the PR description…
self.rearrange is now a PatchTokenizer, not an einops Rearrange, so the attribute name is…
Update migration status to reflect implemented work

        X_patch: Tensor
            [batch, n_chans, n_times//patch_size, patch_size]
        """
+        x = self.tokenizer._prepare_input(x)


                "4D (batch, channels, n_patches, patch_size)."
            )

-        if n_times % self.patch_size != 0:
-            raise ValueError(
-                f"n_times ({n_times}) must be divisible by patch_size ({self.patch_size})."
-            )
        if n_channels % self.in_channels != 0:
            raise ValueError(
                "The input channel dimension "


+        x_signal = self.tokenizer._prepare_input(x_signal)
+        if mask is not None:
+            mask = self.tokenizer._prepare_input(mask)


+        # Shared tokenizer: (batch, n_chans, n_times) -> (batch, n_chans, n_patch, patch_size),
+        # padding/cropping a non-divisible time axis at forward time.
+        self.rearrange = PatchTokenizer(
            patch_size=patch_size,
+            n_times=self._n_times if self._n_times is not None else patch_size,
+            on_non_divisible=on_non_divisible,
        )


codecov · 2026-10-05T12:39:31Z

Codecov Report

❌ Patch coverage is 91.30435% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.99%. Comparing base (2426ddc) to head (16caba1).
⚠️ Report is 7 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1226      +/-   ##
==========================================
+ Coverage   87.84%   87.99%   +0.14%     
==========================================
  Files         151      151              
  Lines       17602    17628      +26     
==========================================
+ Hits        15462    15511      +49     
+ Misses       2140     2117      -23

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Conflict resolution: - braindecode/models/cbramod.py: kept master's braindecode#1226 PatchTokenizer on_non_divisible padding and merged the PR's _knows_geometry() helper (superset of master's _n_times/_n_chans-not-None check) into the PatchTokenizer n_times arg, final_layer construction (both __init__ sites), reset_head, and _n_patch. - docs/whats_new.rst: kept both the braindecode#1233 CBraMod lazy-head bug-fix entry and master's braindecode#1207/braindecode#1212 Deep4Net/ShallowFBCSPNet bug-fix entries. - test/unit_tests/models/test_foundation_models.py: kept master's MAPA/PopT test blocks and the PR's test_cbramod_head_is_concrete_when_geometry_is_derived.

…fig (stacked on #1226) (#1232) * Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraMod, LUNA; ZUNA default on_non_divisible='pad' * from_pretrained: take input geometry from the caller, not the Hub config * ponytail: trim from_pretrained comment, inline one-use test helper

…ights (stacked on #1226) (#1228) * Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraMod, LUNA; ZUNA default on_non_divisible='pad' * Add a geometry-compatibility contract test for every model with released weights * TST stop muting MNE warnings in the geometry contract test mne.set_log_level('ERROR') at import silenced mne.utils.warn for the whole pytest worker, so 15 pytest.warns tests in test_models.py failed with DID NOT WARN when collected after this file. * TST simplify geometry contract test helpers (ponytail) Inline the 19-line TEN_TWENTY list as one split(), and fold the NOT_YET branch of test_geometry_contract into the forward branch (it ran the same assertion). Parametrize ids and strict xfail markers are unchanged.

…ined works (stacked on #1226) (#1233) * Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraMod, LUNA; ZUNA default on_non_divisible='pad' * CBraMod: concrete head whenever the geometry is known (chs_info / input_window_seconds) * CBraMod: tokenizer needs only n_times, keep its divisibility check without channels

…rs on_non_divisible (#1250) * FIX from_pretrained with explicit None geometry; LaBraM decoder honours on_non_divisible - from_pretrained(d, chs_info=None) / (d, n_chans=None) load the saved geometry again: the pinning checks the value, not the key (#1232). - Labram(neural_tokenizer=False) pads / crops / raises on a window not divisible by patch_size, like the tokenizer mode (#1226). - PatchTokenizer.prepare_input is public (_prepare_input kept as alias); LaBraM and LUNA call it. * MAINT PatchTokenizer: drop the unused _prepare_input alias * DOC whats_new entry for #1250 * FIX from_pretrained: drop explicit None geometry so config.json fills it

Copilot AI balanced review requested due to automatic review settings October 5, 2026 10:43

bruAristimunha added a commit to bruAristimunha/braindecode that referenced this pull request Oct 5, 2026

whats_new: reference braindecode#1226

fbc0d62

Copilot started reviewing on behalf of bruAristimunha October 5, 2026 10:44 View session

Copilot AI reviewed Oct 5, 2026

View reviewed changes

bruAristimunha force-pushed the feat/pretrained-compat-d branch from fbc0d62 to 2fd78aa Compare October 5, 2026 11:16

Copilot AI balanced review requested due to automatic review settings October 5, 2026 11:16

This was referenced Oct 5, 2026

Pretrained models compatible by default: input-geometry contract, audits of the 19 pretrained classes, stacked plan #1227

Open

Geometry-compatibility contract test for every model with released weights (stacked on #1226) #1228

Merged

bruAristimunha added the maintenance Bug fix / refactor / tests — not a new model label Oct 5, 2026

Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraM…

16caba1

…od, LUNA; ZUNA default on_non_divisible='pad'

bruAristimunha force-pushed the feat/pretrained-compat-d branch from 2fd78aa to 16caba1 Compare October 5, 2026 11:20

Copilot AI reviewed Oct 5, 2026

View reviewed changes

Copilot AI balanced review requested due to automatic review settings October 5, 2026 11:34

Copilot started reviewing on behalf of bruAristimunha October 5, 2026 11:39 View session

Copilot AI reviewed Oct 5, 2026

View reviewed changes

Copilot started reviewing on behalf of bruAristimunha October 5, 2026 12:03 View session

This was referenced Oct 5, 2026

from_pretrained: take input geometry from the caller, not the Hub config (stacked on #1226) #1232

Merged

CBraMod: concrete head whenever the geometry is known, so from_pretrained works (stacked on #1226) #1233

Merged

bruAristimunha merged commit 5e00a5b into braindecode:master Oct 5, 2026
13 checks passed

bruAristimunha deleted the feat/pretrained-compat-d branch October 5, 2026 15:14

bruAristimunha mentioned this pull request Oct 7, 2026

FIX from_pretrained with explicit None geometry; LaBraM decoder honours on_non_divisible #1250

Merged

Navigation Menu

Repository navigation

Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraMod, LUNA; ZUNA default to pad - #1226

Pad non-divisible windows by default: PatchTokenizer in LaBraM, CBraMod, LUNA; ZUNA default to pad#1226
bruAristimunha merged 1 commit into
braindecode:masterfrom
bruAristimunha:feat/pretrained-compat-d

bruAristimunha commented Oct 5, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

bruAristimunha commented Oct 5, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

codecov Bot commented Oct 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		Status: proposal, branch `feat/pretrained-compat-d`, nothing merged. Companion audit:
		`scripts/compat_audit.py` (geometry grid × 19 pretrained classes, random weights).

Conversation

bruAristimunha commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Testing

Notes for reviewers

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

bruAristimunha commented Oct 5, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

Uh oh!

codecov Bot commented Oct 5, 2026

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bruAristimunha commented Oct 5, 2026 •

edited

Loading