Repository navigation
Geometry-compatibility contract test for every model with released weights (stacked on #1226) - #1228
Merged
bruAristimunha merged 6 commits intoOct 6, 2026
Conversation
3 tasks done
…od, LUNA; ZUNA default on_non_divisible='pad'
bruAristimunha
force-pushed
the
feat/pretrained-compat-test
branch
from
October 5, 2026 11:20
05b5186 to
aba4cf1
Compare
bruAristimunha
force-pushed
the
feat/pretrained-compat-test
branch
from
October 5, 2026 11:21
aba4cf1 to
5403186
Compare
Contributor
There was a problem hiding this comment.
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
Open (8)
PatchTokenizer._prepare_input, which is a private method (leading underscore) and… · Newstate_dict()includes buffers as well as parameters; a… · Newmne.set_log_level(\"ERROR\")sets a global logging level at import time, which can affect… · Newchs_infoentries typically storekindas a FIFF integer constant (as produced by… · Newchs_infoentries typically storekindas a FIFF integer constant (as produced by… · Newchs_infoentries typically storekindas a FIFF integer constant (as produced by… · Newbaseis a mutable list of dicts that’s shared across multiple parametrized cases… · Newself.rearrangepreviously referred to an einopsRearrangelayer but is now aPatchTokenizer.… · New
What changed in this PR
Adds a comprehensive, declarative “geometry compatibility” contract test for all models with released weights, and updates patch-tokenization behavior to support non-divisible windows via padding by default.
Changes:
- Added
test_pretrained_compat.pyto exercise every pretrained-weight model across a grid of channel geometries and window lengths (with strictxfailfor not-yet-supported cells). - Updated foundation-model tests and docs to reflect default padding for non-divisible windows (
on_non_divisible="pad"). - Integrated
PatchTokenizerinto multiple models (e.g., Labram/LUNA/CBraMod) and updated ZUNA defaults/docs accordingly.
| File | Description |
|---|---|
| test/unit_tests/models/test_pretrained_compat.py | New contract test covering model geometry/window compatibility across many scenarios. |
| test/unit_tests/models/test_foundation_models.py | Updates/extends tests to validate default padding behavior and warning/override behavior. |
| docs/whats_new.rst | Documents the new compatibility test and the on_non_divisible default/behavior changes. |
| braindecode/models/zuna.py | Changes default on_non_divisible to "pad" and updates docstring accordingly. |
| braindecode/models/luna.py | Introduces PatchTokenizer to support non-divisible windows; updates patch embedding & token prep. |
| braindecode/models/labram.py | Adds PatchTokenizer-backed handling for non-divisible windows; removes divisibility check from _PatchEmbed. |
| braindecode/models/cbramod.py | Replaces einops Rearrange patching with PatchTokenizer; updates head sizing for padded case. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+284
to
+286
| x_signal = self.tokenizer._prepare_input(x_signal) | ||
| if mask is not None: | ||
| mask = self.tokenizer._prepare_input(mask) |
| 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()) |
Comment on lines
+86
to
+107
| chs = [] | ||
| for n in names: | ||
| loc = np.zeros(12) | ||
| key = upper.get(n.upper()) | ||
| if key is not None: | ||
| loc[:3] = pos[key] | ||
| chs.append({"ch_name": n, "kind": kind, "loc": loc}) | ||
| return chs | ||
|
|
||
|
|
||
| def chs_coords_only(n, kind="eeg"): | ||
| 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": kind, "loc": loc}) | ||
| return chs | ||
|
|
||
|
|
||
| def chs_names_no_loc(names): | ||
| return [{"ch_name": n, "kind": "eeg", "loc": np.zeros(12)} for n in names] |
Comment on lines
+83
to
+107
| def chs_from_montage(names, montage="standard_1005", kind="eeg"): | ||
| pos = _montage(montage).get_positions()["ch_pos"] | ||
| upper = {k.upper(): k for k in pos} | ||
| chs = [] | ||
| for n in names: | ||
| loc = np.zeros(12) | ||
| key = upper.get(n.upper()) | ||
| if key is not None: | ||
| loc[:3] = pos[key] | ||
| chs.append({"ch_name": n, "kind": kind, "loc": loc}) | ||
| return chs | ||
|
|
||
|
|
||
| def chs_coords_only(n, kind="eeg"): | ||
| 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": kind, "loc": loc}) | ||
| return chs | ||
|
|
||
|
|
||
| def chs_names_no_loc(names): | ||
| return [{"ch_name": n, "kind": "eeg", "loc": np.zeros(12)} for n in names] |
Comment on lines
+86
to
+107
| chs = [] | ||
| for n in names: | ||
| loc = np.zeros(12) | ||
| key = upper.get(n.upper()) | ||
| if key is not None: | ||
| loc[:3] = pos[key] | ||
| chs.append({"ch_name": n, "kind": kind, "loc": loc}) | ||
| return chs | ||
|
|
||
|
|
||
| def chs_coords_only(n, kind="eeg"): | ||
| 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": kind, "loc": loc}) | ||
| return chs | ||
|
|
||
|
|
||
| def chs_names_no_loc(names): | ||
| return [{"ch_name": n, "kind": "eeg", "loc": np.zeros(12)} for n in names] |
Comment on lines
+307
to
+311
| base = geos["G1"]["chs_info"] | ||
| 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) |
Comment on lines
+192
to
+196
| # 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, |
Contributor
There was a problem hiding this comment.
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: 2
Open (14)
_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()… · New_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()… · New_SegmentPatch.forward()still statesn_times//patch_size, but… · Newchs_infois a mutable list of dicts, and the same object is reused across multiple geometries… · Newbaseis a mutable list of dicts that’s shared across multiple parametrized cases…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…mne.set_log_level(\"ERROR\")sets a global logging level at import time, which can affect…state_dict()includes buffers as well as parameters; a…PatchTokenizer._prepare_input, which is a private method (leading underscore) and…self.rearrangenow holds aPatchTokenizer, not an einopsRearrangelayer. This name is… · Newself.rearrangepreviously referred to an einopsRearrangelayer but is now aPatchTokenizer.…
Comment on lines
+846
to
+854
| def __init__( | ||
| self, embed_dim: int = 64, patch_size: int = 40, on_non_divisible: str = "pad" | ||
| ) -> None: | ||
| super().__init__() | ||
| self.patch_size = patch_size | ||
| self.embed_dim = embed_dim | ||
| self.tokenizer = PatchTokenizer( | ||
| patch_size=patch_size, n_times=patch_size, on_non_divisible=on_non_divisible | ||
| ) |
| output: (B, C*S, D) where S = T//patch_size, D = embed_dim | ||
| """ | ||
| x = rearrange(x, "B C (S P) -> B (C S) P", P=self.patch_size) | ||
| x = rearrange(self.tokenizer(x), "B C S P -> B (C S) P") |
| Returns: | ||
| -------- | ||
| X_patch: Tensor | ||
| [batch, n_chans, n_times//patch_size, patch_size] |
Comment on lines
+307
to
+312
| base = geos["G1"]["chs_info"] | ||
| 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) | ||
| return geos |
Comment on lines
+190
to
+197
| # 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, | ||
| ) | ||
| self._on_non_divisible = on_non_divisible |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
REVE is skipped, sampling-rate coverage is absent, and some rejection checks can report false positives.
Review effort: Balanced
Findings: 2
Open (17)
_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…chs_infois a mutable list of dicts, and the same object is reused across multiple geometries…_SegmentPatch.forward()still statesn_times//patch_size, but…baseis a mutable list of dicts that’s shared across multiple parametrized cases…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…mne.set_log_level(\"ERROR\")sets a global logging level at import time, which can affect…state_dict()includes buffers as well as parameters; a…PatchTokenizer._prepare_input, which is a private method (leading underscore) and…self.rearrangenow holds aPatchTokenizer, not an einopsRearrangelayer. This name is…self.rearrangepreviously referred to an einopsRearrangelayer but is now aPatchTokenizer.…
Comment on lines
+270
to
+281
| ("BENDR", "G4"), | ||
| ("BENDR", "G2"), | ||
| ("BENDR", "G3"), | ||
| ("BENDR", "G3b"), # ChannelTokenizer(fixed_order) | ||
| ( | ||
| "EEGDINO", | ||
| "G2", | ||
| ), # ChannelTokenizer(index_slots): > 19 channels needs a declared error | ||
| ( | ||
| "SignalJEPA", | ||
| "G3b", | ||
| ), # names without coordinates: output is NaN (division by zero) today |
| strict=True, reason="not migrated yet (see design doc)" | ||
| ) | ||
| ) | ||
| yield pytest.param(name, gname, gkw, id=f"{name}-{gname}", marks=marks) |
Comment on lines
+309
to
+311
| 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) |
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.
Contributor
There was a problem hiding this comment.
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: 2
Open (20)
_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…baseis the same mutablechs_infolist reused acrossG1/G5a/G5b/G5c. If any model… · Newwarnings.simplefilter('ignore')suppresses all warnings from model construction/forward, which… · Newwant == 'raise', the test currently accepts anyValueError/RuntimeError, which can allow… · Newchs_infois a mutable list of dicts, and the same object is reused across multiple geometries…_SegmentPatch.forward()still statesn_times//patch_size, but…baseis a mutable list of dicts that’s shared across multiple parametrized cases…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…mne.set_log_level(\"ERROR\")sets a global logging level at import time, which can affect…state_dict()includes buffers as well as parameters; a…PatchTokenizer._prepare_input, which is a private method (leading underscore) and…self.rearrangenow holds aPatchTokenizer, not an einopsRearrangelayer. This name is…self.rearrangepreviously referred to an einopsRearrangelayer but is now aPatchTokenizer.…
Comment on lines
+305
to
+309
| base = geos["G1"]["chs_info"] | ||
| 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) |
Comment on lines
+359
to
+363
| with warnings.catch_warnings(): | ||
| warnings.simplefilter("ignore") | ||
| model = spec["cls"](**kw).eval() | ||
| with torch.no_grad(): | ||
| return model(torch.randn(1, n_ch, gkw["n_times"])) |
Comment on lines
+371
to
+373
| if want == "raise": | ||
| with pytest.raises((ValueError, RuntimeError)): | ||
| build_and_forward() |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #1228 +/- ##
==========================================
+ Coverage 88.41% 88.42% +0.01%
==========================================
Files 156 156
Lines 19009 19009
==========================================
+ Hits 16806 16809 +3
+ Misses 2203 2200 -3 🚀 New features to boost your workflow:
|
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.
Contributor
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
EEGDINO asserts the wrong xfail target, REVE is skipped by default, and sampling-rate geometries are not generated.
Review effort: Balanced
Findings: 2
Open (19)
_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…want == 'raise', the test currently accepts anyValueError/RuntimeError, which can allow…warnings.simplefilter('ignore')suppresses all warnings from model construction/forward, which…baseis the same mutablechs_infolist reused acrossG1/G5a/G5b/G5c. If any model…chs_infois a mutable list of dicts, and the same object is reused across multiple geometries…_SegmentPatch.forward()still statesn_times//patch_size, but…baseis a mutable list of dicts that’s shared across multiple parametrized cases…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…mne.set_log_level(\"ERROR\")sets a global logging level at import time, which can affect…state_dict()includes buffers as well as parameters; a…PatchTokenizer._prepare_input, which is a private method (leading underscore) and…self.rearrangenow holds aPatchTokenizer, not an einopsRearrangelayer. This name is…self.rearrangepreviously referred to an einopsRearrangelayer but is now aPatchTokenizer.…
Resolved since last review (1)
Resolved docs/whats_new.rst by keeping both entries.
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The EEGDINO expectation is inverted, and released-weight models are missing from the advertised comprehensive contract.
Review effort: Balanced
Findings: 2
Open (20)
_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…_PatchEmbedNetworkinitializesPatchTokenizerwithn_times=patch_size, butforward()…want == 'raise', the test currently accepts anyValueError/RuntimeError, which can allow…warnings.simplefilter('ignore')suppresses all warnings from model construction/forward, which…baseis the same mutablechs_infolist reused acrossG1/G5a/G5b/G5c. If any model…chs_infois a mutable list of dicts, and the same object is reused across multiple geometries…_SegmentPatch.forward()still statesn_times//patch_size, but…baseis a mutable list of dicts that’s shared across multiple parametrized cases…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…chs_infoentries typically storekindas a FIFF integer constant (as produced by…state_dict()includes buffers as well as parameters; a…PatchTokenizer._prepare_input, which is a private method (leading underscore) and…self.rearrangenow holds aPatchTokenizer, not an einopsRearrangelayer. This name is…self.rearrangepreviously referred to an einopsRearrangelayer but is now aPatchTokenizer.…
Resolved since last review (1)
| channels="coords", | ||
| coords_checked=False, | ||
| ), | ||
| } |
Comment on lines
+6
to
+7
| names without coordinates, short / long / non-divisible windows, other | ||
| sampling rates -- and must either forward or raise the *declared* error. |
bruAristimunha
added a commit
to bruAristimunha/braindecode
that referenced
this pull request
Oct 6, 2026
test_pretrained_compat.py (add/add after braindecode#1228's squash merge): keep this PR's side, i.e. braindecode#1228's file with the six NOT_YET cells and the fixed_order raise removed, since ChannelTokenizer makes them pass. whats_new: keep both; drop the EEGDINO clause (this PR does not change EEGDINO).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacked on #1226 (review the last commit only). Part of #1227.
One test that checks every model with released weights on a grid of input geometries — the mechanism timm uses (
test_model_load_pretrainedoverlist_models(pretrained=True)), applied to EEG's axes: canonical montage, permuted order, a 64-channel montage outside the 10-20 vocabulary, coordinates-only channels, names without coordinates, 1 s / 30 s / non-divisible windows. Random weights, no download.Changes
test/unit_tests/models/test_pretrained_compat.py: 19 classes × 9 geometries. Expected outcome per cell is derived from one declared channel strategy per model (COMPAT), not hard-coded; cells not handled yet arexfail(strict=True)(LaBraM outside its vocabulary, BENDR on any other order, EEGDINO > 19 channels, SignalJEPA NaN with names-without-coords), so each fix flips a marker.docs/whats_new.rst.Testing
pytest test/unit_tests/models/test_pretrained_compat.py→ 128 passed, 8 skipped (REVE: network-marked by conftest), 9 xfailed (~19 s on CPU).ruffclean.Regression test (this is one)
Style checks recorded
docs/whats_new.rstupdatedNotes for reviewers
Found while writing it: SignalJEPA returns NaN when channel names have zero
loc(division by zero in the channel embedding) — tracked in #1227, not fixed here.