BrainOmni: rotate with the released checkpoint's RoPE cache; keep the released head init - #1244

What changes

RoPE. The released BrainOmni checkpoints (OpenTSLab/BrainOmni @ 9a4d3c70, tiny and base) store blocks.*.time_attn.rope_embedding_layer.rotate as a float32 tensor of shape (240, n_dim/2). It holds the cosines cos(t·freqs) (max-abs 6e-8 from the exact values). The released downstream loader (downstream/model_collection/BrainOmni.py, load_state_dict(strict=False)) copies this tensor into the complex64 rotate buffer of model_utils/attn.py::RotaryEmbedding. The sine part is therefore 0, and the released code rotates queries and keys with that cache. Every downstream sequence is shorter than 240 (Somato 2 s → 8 tokens), so all of them use the cached rotation. freqs is stored rounded to bfloat16 and is used only past 240 positions.

The port dropped both keys and recomputed cos and sin. With this PR:

• _RotaryPositionalEmbedding has the release's freqs and rotate buffers. rotate is stored as the real view (seq, n_dim/2, 2) = (cos, sin), so it still runs in real arithmetic on devices without complex support.
• The official keys are loaded. A real (seq, n_dim/2) cache is read as (cos, 0), which is what the release's copy_ into a complex buffer does.
• Positions past the cache use freqs with both cos and sin, as in the release's _set_rotate_cache.
• The buffers stay float32 under .half()/.bfloat16(), like the release's complex64 cache under a bf16 engine.
• Native state dicts saved without the buffers still load.

Head init. final_layer keeps PyTorch's default Linear initialisation. The released head (downstream/model.py: Dropout → LazyLinear → SELU → Linear) is never passed through _init_weights, which brainomni/model.py applies only to the backbone.

Numbers

Forward parity on identical inputs and identical released weights: 8 test segments of SomatoMotor (ds006035, 372 channels × 512 samples) prepared by the released maker. The released code is run verbatim (complex RoPE, CPU fp32) with attention dropout set to 0, because the release applies SDPA dropout in eval as well. Values are the max-abs difference per layer output:

All 11 blocks used by encode are 0 with this PR.

In bf16 (both models cast with .bfloat16()), encode differs from the released code by 0.0586. The released code's own bf16 output differs from its fp32 output by 0.0597.

Fine-tuning on SomatoMotor (372-channel input, BrainOmni tiny, the released downstream protocol on the same prepared segments and folds: 3 LRs × 5 subject folds × 2 seeds, best-val checkpoint tested, mean ± sd of test balanced accuracy at the LR with the best test mean, as in the released metrics_stat.py):

The paper reports 0.863 ± 0.128 for this cell. A full grid with this branch is running.

Not changed: the released attention applies dropout_p in eval mode too. The port keeps eval deterministic (row "eval-mode attention dropout" above).

Tests

• test_spatial_temporal_block_matches_frozen_released_output: a seeded block whose RoPE buffers are in the released format, checked against values from the released SpatialTemporalAttentionBlock on the same state dict and input (offline).
• test_rope_rotates_with_the_released_cosine_only_cache (5 / 240 / 250 positions), test_rope_buffers_stay_float32_and_keep_dtype, test_brainomni_state_dict_without_rope_buffers_loads, test_brainomni_head_keeps_default_linear_init; the official-key load test now checks the converted cache.
• The network parity test is re-pinned to the released code's output on the same fixture (sum 51.254784; first 16 values max-abs 0.0 against the released code).
• test_brainomni.py: 110 passed (network tests deselected locally). test_models / test_integration / test_return_features / test_model_contract / test_pretrained_compat / test_correctness / test_huggingface, -k brainomni: 19 passed, 1 skipped. pre-commit passes.