Repository navigation
BrainOmni: load braindecode-hosted weights, drop the OpenTSLab translation - #1245
Merged
bruAristimunha merged 9 commits intoOct 8, 2026
Merged
Conversation
…eep the released head init
# Conflicts: # docs/whats_new.rst
… dtype, as the release does
…a partial RoPE state dict
…translation, cut the port
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #1245 +/- ##
==========================================
- Coverage 89.05% 89.00% -0.05%
==========================================
Files 157 157
Lines 19654 19367 -287
==========================================
- Hits 17503 17238 -265
+ Misses 2151 2129 -22 🚀 New features to boost your workflow:
|
# Conflicts: # braindecode/models/brainomni.py # docs/whats_new.rst # test/unit_tests/models/test_brainomni.py
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Window stride and attention-dimension validation regressions can produce division-by-zero failures for public API inputs.
2 open findings
What changed in this PR
Migrates BrainOmni models to native Braindecode-hosted checkpoints while simplifying checkpoint handling and architecture code.
Changes:
- Uses standard
from_pretrainedloading without OpenTSLab translation. - Refactors windowing, attention, SEANet, and weight normalization.
- Updates compatibility, parity, integration tests, and documentation.
| File | Description |
|---|---|
braindecode/models/brainomni.py |
Simplifies model architecture and checkpoint loading. |
braindecode/modules/quantization.py |
Removes the fixed EMA strategy argument. |
test/unit_tests/models/test_brainomni.py |
Updates focused and pretrained parity tests. |
test/unit_tests/models/test_integration.py |
Exempts the tokenizer’s fixed activation. |
test/unit_tests/models/test_pretrained_compat.py |
Adds BrainOmni compatibility coverage. |
docs/api.rst |
Updates pretrained-weight availability. |
docs/whats_new.rst |
Documents converted Hub checkpoints. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
# Conflicts: # test/unit_tests/models/test_pretrained_compat.py
bruAristimunha
added a commit
to bruAristimunha/braindecode
that referenced
this pull request
Oct 8, 2026
…aindecode#1250, braindecode#1252) into fix/hpu-models
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.
Stacked on #1244 (base commit 517aa14); the diff below is against that commit.
What changed
BrainOmni and BrainTokenizer load with the standard
from_pretrained. The weights are converted once from OpenTSLab/BrainOmni@9a4d3c70 (MIT):braindecode/braintokenizer-pretrained(rev 945aa489)braindecode/brainomni-tiny-pretrained(rev d111525e)braindecode/brainomni-base-pretrained(rev 40f33629)Each repo has
config.json,model.safetensors,pytorch_model.bin,README.mdandconvert_brainomni_checkpoints.py.Removed:
BrainTokenizer.from_opentslab_config,BrainOmni.from_opentslab_config, the twoload_state_dictoverrides,_rename_official_key,_translate_opentslab_config,_TOKENIZER_CONFIG_RENAMES,_BRAINOMNI_CONFIG_RENAMESand_BRAINOMNI_PRETRAINING_KEYS. No key mapping is left.Removed arguments:
quantize_optimize_method(BrainTokenizer, BrainOmni, ResidualVectorQuantizer) andactivation(BrainTokenizer, which never used it).State dict:
posandsensor_typeare now non-persistent, because they are derived fromchs_info.weight_g/weight_vbecomeparametrizations.weight.original0/1, sincetorch.nn.utils.weight_normis deprecated..model.level is flattened.Windowing is now 3 lines with
Tensor.unfold._unfoldand_window_strideare removed.The three attention modules share one SDPA helper. SEANet keeps only the released configuration.
Tests:
test_brainomni.pynow loads the new repo ids. BrainOmni is added totest_pretrained_compat.Parity (fingerprint of 91 entries: tiny, base and tokenizer weights in fp32/bf16/fp16; T = 512, 1024, 7680 and 12032; random-init outputs, gradients and buffers)
Lines
braindecode/models/brainomni.pytest/unit_tests/models/test_brainomni.py