Repository navigation
CBraMod: concrete head whenever the geometry is known, so from_pretrained works (stacked on #1226) - #1233
Merged
bruAristimunha merged 5 commits intoOct 6, 2026
Conversation
…od, LUNA; ZUNA default on_non_divisible='pad'
…ut_window_seconds)
bruAristimunha
force-pushed
the
fix/cbramod-lazy-head
branch
from
October 5, 2026 13:16
d2e483b to
cf66b6c
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: 3
Open (4)
_prepare_inputis a private method ofPatchTokenizer, so calling it from model code tightly… · New_PatchEmbedNetworkbuilds its ownPatchTokenizerconfigured withn_times=patch_size, which is… · Newnot any('tokenizer' in k for k in model.state_dict())is string-based and can miss… · Newrearrangeis now aPatchTokenizerthat may pad/crop and tokenize, not just… · New
What changed in this PR
Fixes CBraMod’s from_pretrained failure by ensuring the classification head is concrete (Linear) whenever channel/time geometry is known (including when derived from chs_info / input_window_seconds), while standardizing non-divisible window handling via PatchTokenizer across models.
Changes:
- CBraMod: detect “known geometry” via public properties and build
Linearhead accordingly; keepLazyLinearonly for truly unknown geometry. - Labram/LUNA/ZUNA: route time-axis tokenization through
PatchTokenizerwith newon_non_divisiblebehavior (ZUNA default now"pad"). - Tests + docs: add regression coverage for padding/default behavior and for CBraMod save/load round-trip; document behavior change and bug fix.
| File | Description |
|---|---|
| test/unit_tests/models/test_foundation_models.py | Updates/extends tests for default padding behavior and CBraMod save/load regression. |
| docs/whats_new.rst | Documents non-divisible window behavior changes and CBraMod LazyLinear head bug fix. |
| braindecode/models/zuna.py | Changes on_non_divisible default to "pad" and updates docstring accordingly. |
| braindecode/models/luna.py | Introduces PatchTokenizer for time tokenization and adds on_non_divisible parameter. |
| braindecode/models/labram.py | Adds on_non_divisible support and centralizes pad/crop/error handling via PatchTokenizer. |
| braindecode/models/cbramod.py | Uses PatchTokenizer + “known geometry” logic to avoid non-serializable lazy heads when geometry is derivable. |
💡 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) |
Comment on lines
+847
to
+854
| 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 | ||
| ) |
Comment on lines
+2004
to
+2016
| @pytest.mark.parametrize("cls,kwargs,patch_size", _patch_models()) | ||
| def test_non_divisible_window_padded_by_default(cls, kwargs, patch_size): | ||
| """A window that is not a multiple of patch_size pads (warning) and only | ||
| raises with on_non_divisible="error"; the tokenizer adds no parameters.""" | ||
| kwargs = dict(kwargs, n_times=kwargs["n_times"] + 37, n_outputs=2) | ||
| with pytest.warns(UserWarning, match="not divisible"): | ||
| model = cls(**kwargs).eval() | ||
| n_chans = kwargs.get("n_chans") or len(kwargs["chs_info"]) | ||
| x = torch.randn(1, n_chans, kwargs["n_times"]) | ||
| 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
+190
to
196
| # 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._knows_geometry() 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: 1
Open (9)
xviaPatchTokenizer, butmaskis not passed through the same… · New_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes… · New_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes… · NewPatchTokenizer._prepare_input, which is a private method (leading underscore)… · Newnot any('tokenizer' in k for k in model.state_dict())is string-based and can miss…_PatchEmbedNetworkbuilds its ownPatchTokenizerconfigured withn_times=patch_size, which is…_prepare_inputis a private method ofPatchTokenizer, so calling it from model code tightly…rearrangeis now aPatchTokenizerthat may pad/crop and tokenize, not just…
Comment on lines
+190
to
196
| # 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._knows_geometry() else patch_size, | ||
| on_non_divisible=on_non_divisible, | ||
| ) |
Comment on lines
+1035
to
+1037
| self.tokenizer = PatchTokenizer( | ||
| patch_size=patch_size, n_times=n_times, on_non_divisible=on_non_divisible | ||
| ) |
| X_patch: Tensor | ||
| [batch, n_chans, n_times//patch_size, patch_size] | ||
| """ | ||
| x = self.tokenizer._prepare_input(x) |
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) |
| :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 |
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.
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: 1
Open (14)
xviaPatchTokenizer, butmaskis not passed through the same…PatchTokenizer._prepare_input, which is a private method (leading underscore)…_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes…_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes…not any('tokenizer' in k for k in model.state_dict())is string-based and can miss…_PatchEmbedNetworkbuilds its ownPatchTokenizerconfigured withn_times=patch_size, which is…_prepare_inputis a private method ofPatchTokenizer, so calling it from model code tightly…_knows_geometry()is called multiple times during initialization/reset, which repeats the same… · New_knows_geometry()is called multiple times during initialization/reset, which repeats the same… · New_knows_geometry()is called multiple times during initialization/reset, which repeats the same… · Newself.n_chans, self.n_timessolely for its side-effect (property access) is a bit opaque.… · Newtorch.no_grad()context avoids… · Newrearrangeis now aPatchTokenizerthat may pad/crop and tokenize, not just…
| self.rearrange = PatchTokenizer( | ||
| patch_size=patch_size, | ||
| n_times=self._n_times if self._n_times is not None else patch_size, | ||
| n_times=self.n_times if self._knows_geometry() else patch_size, |
| if return_encoder_output: | ||
| self.final_layer = nn.Identity() | ||
| elif self._n_times is not None and self._n_chans is not None: | ||
| elif self._knows_geometry(): |
| # A head implies a classifier, also when built with return_encoder_output. | ||
| self._update_init_kwargs(return_encoder_output=False) | ||
| if self._n_times is not None and self._n_chans is not None: | ||
| if self._knows_geometry(): |
| head can be a real Linear instead of a LazyLinear that cannot be saved | ||
| or loaded before a forward pass.""" | ||
| try: | ||
| self.n_chans, self.n_times |
| model.save_pretrained(tmp_path) | ||
| loaded = CBraMod.from_pretrained(tmp_path) | ||
| x = torch.randn(1, len(chs), 800) | ||
| assert torch.allclose(model.eval()(x), loaded.eval()(x), atol=1e-5) |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #1233 +/- ##
==========================================
+ Coverage 88.40% 88.41% +0.01%
==========================================
Files 156 156
Lines 18999 19019 +20
==========================================
+ Hits 16796 16816 +20
Misses 2203 2203 🚀 New features to boost your workflow:
|
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: 1
Open (16)
xviaPatchTokenizer, butmaskis not passed through the same…PatchTokenizer._prepare_input, which is a private method (leading underscore)…_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes…_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes…not any('tokenizer' in k for k in model.state_dict())is string-based and can miss…_PatchEmbedNetworkbuilds its ownPatchTokenizerconfigured withn_times=patch_size, which is…_prepare_inputis a private method ofPatchTokenizer, so calling it from model code tightly…self.n_chans, self.n_timespurely for its side-effects is a bit opaque and can look like a… · Newtorch.testing.assert_close(...)typically provides clearer failure messages… · Newtorch.no_grad()context avoids…self.n_chans, self.n_timessolely for its side-effect (property access) is a bit opaque.…_knows_geometry()is called multiple times during initialization/reset, which repeats the same…_knows_geometry()is called multiple times during initialization/reset, which repeats the same…_knows_geometry()is called multiple times during initialization/reset, which repeats the same…rearrangeis now aPatchTokenizerthat may pad/crop and tokenize, not just…
| head can be a real Linear instead of a LazyLinear that cannot be saved | ||
| or loaded before a forward pass.""" | ||
| try: | ||
| self.n_chans, self.n_times |
| model.save_pretrained(tmp_path) | ||
| loaded = CBraMod.from_pretrained(tmp_path) | ||
| x = torch.randn(1, len(chs), 800) | ||
| assert torch.allclose(model.eval()(x), loaded.eval()(x), atol=1e-5) |
Resolved docs/whats_new.rst by keeping both entries.
Contributor
There was a problem hiding this comment.
Copilot review overview
🟢 Approval recommended
The implementation consistently uses resolved geometry and includes focused regression coverage for serialization and fallback behavior.
Review effort: Balanced
Findings: 1
Open (16)
xviaPatchTokenizer, butmaskis not passed through the same…PatchTokenizer._prepare_input, which is a private method (leading underscore)…_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes…_SegmentPatchdepends on the privatePatchTokenizer._prepare_input. This makes…not any('tokenizer' in k for k in model.state_dict())is string-based and can miss…_PatchEmbedNetworkbuilds its ownPatchTokenizerconfigured withn_times=patch_size, which is…_prepare_inputis a private method ofPatchTokenizer, so calling it from model code tightly…torch.testing.assert_close(...)typically provides clearer failure messages…self.n_chans, self.n_timespurely for its side-effects is a bit opaque and can look like a…torch.no_grad()context avoids…self.n_chans, self.n_timessolely for its side-effect (property access) is a bit opaque.…_knows_geometry()is called multiple times during initialization/reset, which repeats the same…_knows_geometry()is called multiple times during initialization/reset, which repeats the same…_knows_geometry()is called multiple times during initialization/reset, which repeats the same…rearrangeis now aPatchTokenizerthat may pad/crop and tokenize, not just…
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
CBraMod.from_pretrained("braindecode/cbramod-pretrained", chs_info=…, n_times=…, sfreq=…)fails at every geometry withAttempted to use an uninitialized parameter … LazyModule(audit in #1227, 10/10 cells).Root cause: CBraMod sizes its head from the private
_n_chans/_n_times, which stayNonewhen the geometry is given aschs_info/input_window_seconds; the model then builds aLazyLinearhead althoughn_chansandn_timesare perfectly known, and aLazyLinearcannot be loaded (or saved) before a forward pass.Stacked on #1226 (review the last commit). Part of #1227.
Changes
cbramod.py:_knows_geometry()(the public properties resolve) decides betweenLinearandLazyLinear; head and_n_patch()useself.n_chans/self.n_times. Lazy fallback kept for a truly unknown geometry.test_foundation_models.py: geometry fromchs_info+input_window_seconds→ concrete head,save_pretrained→from_pretrainedround trip equal, lazy head only when geometry is unknown.docs/whats_new.rst(Bug fixes).Testing
pytest test/unit_tests/models/test_foundation_models.py test/unit_tests/models/test_integration.py -k cbramod→ 14 passed, 2 skipped.pre-commit run --files <changed>clean. Not run: full suite (CI); the Hub checkpoint itself (cluster audit will be re-run after the stack lands).Regression test
Style checks recorded
docs/whats_new.rstupdated