Fix vmap axes for input_positions in dot_product_attention_with_rope

[TheSaiEaranti]

What does this PR do?

Fixes the input_positions argument of dot_product_attention_with_rope (added in #5481), which crashes on every call that passes it:

ValueError: vmap got inconsistent sizes for array axes to be mapped:
  * one axis had size 2: axis -2 of argument x of type float32[1,6,2,8];
  * one axis had size 6: axis -1 of argument input_positions of type float32[6]

The heads-axis vmap used in_axes=(-2, -1), mapping the positions' last axis — which is seq_len, not a heads axis — so any call with num_heads != seq_len raises, and out_axes=(-2, -1) was a 2-tuple for a function returning a single array. Both docstrings (dot_product_attention_with_rope and MultiHeadAttention.__call__) document input_positions of shape (..., seq_len,), and there was no test coverage for it.

Positions are shared by every head, so the fix broadcasts them: jax.vmap(rope, in_axes=(-2, None), out_axes=-2). Since in_axes=None also handles input_positions=None, the single vmap replaces the previous branch.

Tests

test_dot_product_attention_with_rope_input_positions pins three properties:

• explicit arange(seq_len) positions produce the same output as the default (they are the default),
• a uniform shift preserves the output — RoPE encodes relative position, so arange + 3 leaves every pairwise difference unchanged, which distinguishes real broadcasting from positions being ignored per-head,
• non-uniform positions (arange * 2) change the output, proving positions take effect.

The test fails on main (with the vmap error above) and passes with the fix. Full tests/nnx/nn/attention_test.py: 127 passed. ruff check clean.

Checklist

• This change is discussed in a Github issue/discussion — the bug is in a feature merged three days ago (Add RoPE embeddings #5481); filed directly as a fix given the crash is unconditional.
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests.