fix(collator): pass batch index positionally in single-subseq mrope path

[he-yufeng]

What

MultiModalDataCollatorForSeq2Seq._compute_rope_position_ids_with_packing crashes on the single-subsequence branch:

mm_inputs_for_sample = _slice_mm_inputs_for_sample(
    mm_inputs, batch_imglens, batch_vidlens, sample_idx=sample_idx
)

_slice_mm_inputs_for_sample has no sample_idx parameter — its fourth parameter is batch_idx — so any packed multimodal mrope batch that reaches the num_sub_seqs <= 1 path raises:

TypeError: _slice_mm_inputs_for_sample() got an unexpected keyword argument 'sample_idx'

The parallel num_sub_seqs > 1 branch a few lines below already passes the index positionally (_slice_mm_inputs_for_sample(mm_inputs, batch_imglens, batch_vidlens, sample_idx, ...)), so only the single-subsequence branch is affected. Fixes #10497.

Fix

Pass the batch index positionally, matching the packed branch.

Test

Added test_packing_rope_single_subseq_slices_without_crashing to tests/data/test_collator.py. It drives _compute_rope_position_ids_with_packing through the num_sub_seqs <= 1 branch (mocking the per-sample rope call to isolate the slicing). The test raises the TypeError on main and passes with the fix; the existing test_collator.py suite stays green (5 passed).

PYTHONPATH=src python -m pytest tests/data/test_collator.py -q   # 5 passed

[gemini-code-assist]

Code Review

This pull request fixes a keyword argument crash in _compute_rope_position_ids_with_packing by passing sample_idx positionally to _slice_mm_inputs_for_sample, and adds a corresponding regression test. The reviewer identified a bug in the same branch where input_ids is not sliced for the current sample, leading to a shape mismatch when the batch size is greater than 1. Additionally, the reviewer suggested adding the @pytest.mark.runs_on decorator to the new test function for consistency and proper CI/CD execution.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In the single-subsequence branch (num_sub_seqs <= 1), sample_features["input_ids"] is assigned the entire batch's input_ids (features["input_ids"]) instead of being sliced for the current sample (features["input_ids"][sample_idx : sample_idx + 1]).

When the batch size is greater than 1 (bsz > 1), this causes a shape mismatch because input_ids has shape (bsz, seq_len) while attention_mask is correctly sliced to shape (1, seq_len). This leads to a crash in _compute_rope_position_ids or a shape mismatch error when merging the position IDs at the end of the function.

We should slice input_ids to match the current sample, just like it is done in the num_sub_seqs > 1 branch.

Suggested change

	"input_ids": features["input_ids"],
	"input_ids": features["input_ids"][sample_idx : sample_idx + 1],

[gemini-code-assist]

For consistency with the rest of the test suite in this file and to ensure proper execution in the CI/CD pipeline, please add the @pytest.mark.runs_on decorator to the new test function.

@pytest.mark.runs_on(["cpu", "mps"])
def test_packing_rope_single_subseq_slices_without_crashing(monkeypatch):