Honor partial_rotary_factor on the MRoPE path

[adityasingh2400]

Description

partial_rotary_factor was silently dropped on the MRoPE path, so the two shipped configs that set both use_mrope: true and partial_rotary_factor: 0.25 applied RoPE to the full head dimension instead of the leading quarter. Reported in #4616.

The defect: Attention.init_rotary_embedding in src/maxtext/layers/attentions.py routes every use_mrope layer to Qwen3OmniMoeThinkerTextRotaryEmbedding through the elif self.use_mrope: branch. That branch passed embedding_dims=rope_embedding_dims, which is the full head_dim, and never passed a partial rotary factor because the class did not accept one. PartialRotaryEmbedding already implements the intended behavior, but only the non-MRoPE branches can reach it, and the is_qwen3_hybrid branch that does construct it sits below the MRoPE branch in the same chain. There was no error and no shape mismatch, so the misconfiguration was invisible at runtime.

Affected shipped configs, both head_dim: 256:

• src/maxtext/configs/models/qwen3.5-35b-a3b.yml
• src/maxtext/configs/models/qwen3.5-397b-a17b.yml

Both set partial_rotary_factor: 0.25, so 64 of the 256 head channels should be rotated. Their mrope_section of [11, 11, 10] sums to 32, which is half of that intended 64 channel slice rather than half of head_dim, so the frequency layout was wrong too: the interleaved MRoPE sections only reached the first 33 of the 128 angles that the full head dimension produces.

The fix teaches Qwen3OmniMoeThinkerTextRotaryEmbedding the construction PartialRotaryEmbedding already uses. It takes a partial_rotary_factor, computes rotary_dim = int(head_dim * partial_rotary_factor), and initializes the base class with embedding_dims=self.rotary_dim so the timescale, and therefore the inverse frequencies, span only the rotated slice. In __call__ the leading rotary_dim channels are split off, rotated, and the untouched remainder is concatenated back. The shape check now compares against head_dim instead of the reduced embedding_dims. init_rotary_embedding passes config.partial_rotary_factor through, matching what the neighboring qwen3 hybrid branch already does.

The factor defaults to 1.0. In that case no split happens and the output is unchanged, so qwen3-vl-* and qwen3-omni-*, which do not set the factor, are unaffected. I verified this directly: dumping the default layer output at head_dim=128 with mrope_section=(24, 20, 20), for both 3D multimodal positions and 2D text positions, gives bitwise identical arrays before and after the change.

Tests

Five new cases in tests/unit/partial_rotary_embedding_test.py, in a Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest class alongside the existing partial rotary tests:

• channels beyond rotary_dim pass through untouched and the leading slice is rotated,
• with text-only positions MRoPE degenerates to ordinary RoPE, so the output must agree with PartialRotaryEmbedding channel for channel,
• a snapshot of the default full rotation output, which pins the existing behavior,
• an explicit factor of 1.0 matches the default construction,
• the two affected shipped configs are loaded through pyconfig, their mrope_section is checked against rotary_dim // 2, and the passthrough property is asserted at their real head_dim of 256.

Command:

python -m pytest tests/unit/partial_rotary_embedding_test.py -q

Fail before, with src/maxtext/layers/embeddings.py and src/maxtext/layers/attentions.py checked out from the base ref and only the test file carrying the change:

E     TypeError: Qwen3OmniMoeThinkerTextRotaryEmbedding.__init__() got an unexpected keyword argument 'partial_rotary_factor'

=========================== short test summary info ============================
FAILED tests/unit/partial_rotary_embedding_test.py::Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest::test_mrope_full_factor_matches_default
FAILED tests/unit/partial_rotary_embedding_test.py::Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest::test_mrope_partial_matches_partial_rotary_embedding
FAILED tests/unit/partial_rotary_embedding_test.py::Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest::test_mrope_partial_rotary_passthrough
SUBFAILED(model_name='qwen3.5-35b-a3b') tests/unit/partial_rotary_embedding_test.py::Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest::test_shipped_mrope_configs_apply_partial_rotary
SUBFAILED(model_name='qwen3.5-397b-a17b') tests/unit/partial_rotary_embedding_test.py::Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest::test_shipped_mrope_configs_apply_partial_rotary
5 failed, 10 passed in 8.98s

Pass after, with both source files restored:

.............                                                          [100%]
13 passed, 2 subtests passed in 9.37s

The pre-existing cases in that file, including the PartialRotaryEmbedding and Gemma4PartialRotaryEmbedding suites, pass in both runs. tests/unit/embeddings_test.py is also green. Ran on CPU. Formatted with the pinned pyink 24.10.1 at --pyink-indentation=2 --line-length=122, and pylint 3.3.8 rates the changed files 10.00/10.

FIXES: #4616

[adityasingh2400]

The source fix here is superseded by #4764, which merged earlier today. I have rebased and dropped it, so this is now tests only.

Your version is the better one. It adds an instance-level override in attentions.py:

partial_rotary_factor=(
    self.partial_rotary_factor if self.partial_rotary_factor is not None
    else self.config.partial_rotary_factor
),

plus validation of the factor and the rotary dim in MRoPE.__init__. Mine had neither, so there is nothing left worth keeping from my source changes.

What is left is the MRoPE regression coverage, which main does not have. The existing tests/unit/partial_rotary_embedding_test.py covers PartialRotaryEmbeddingTest and Gemma4PartialRotaryEmbeddingTest, but nothing exercises the Qwen3-Omni MRoPE path that #4616 reported. This adds Qwen3OmniMoeThinkerTextRotaryEmbeddingPartialTest with five cases:

• a partial factor leaves the unrotated tail untouched
• MRoPE partial output matches PartialRotaryEmbedding for the same factor
• the default rotates the full head dim
• an explicit factor of 1.0 matches the default
• every shipped MRoPE config actually applies its partial factor, checking rotary_dim and sum(mrope_section) == rotary_dim // 2

Worth saying plainly: I could not run these locally. They were written against my implementation and the repo needs pathwaysutils and the rest of the config chain to import, which I do not have here. They read as compatible with yours, no assertions on error strings and the default is 1.0 either way, but CI is the real check. If any of them are wrong against your version I will fix them.

Close this if the coverage is not worth carrying. The bug itself is fixed either way, and #4616 is still open if you want to close that too.