From f2f0bb7ce56d9571088fb9575d9ca9e36ec6f56d Mon Sep 17 00:00:00 2001 From: SaiEaranti Date: Sun, 26 Jul 2026 15:26:32 -0500 Subject: [PATCH] Fix vmap axes for input_positions in dot_product_attention_with_rope Passing input_positions crashed with "vmap got inconsistent sizes for array axes to be mapped": the heads-axis vmap used in_axes=(-2, -1), mapping the positions' last axis, which is seq_len rather than a heads axis, and out_axes was a 2-tuple for a function returning a single array. Positions are shared by every head, so broadcast them with in_axes=(-2, None). The single vmap also covers the input_positions=None case, replacing the branch. Add a regression test: explicit arange positions match the default, a uniform shift preserves the output (RoPE encodes relative position), and non-uniform positions change it. --- flax/nnx/nn/attention.py | 8 +++----- tests/nnx/nn/attention_test.py | 35 ++++++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/flax/nnx/nn/attention.py b/flax/nnx/nn/attention.py index 025ebdfdd..5bcf887b8 100644 --- a/flax/nnx/nn/attention.py +++ b/flax/nnx/nn/attention.py @@ -469,11 +469,9 @@ def dot_product_attention_with_rope( Output of shape ``[batch..., q_length, num_heads, v_dim]``. """ # query/key: [batch..., seq_length, num_heads, head_dim] - # RoPE.__call__ expects (..., seq_length, head_dim), so vmap over heads. - if input_positions is None: - apply = jax.vmap(rope, in_axes=-2, out_axes=-2) - else: - apply = jax.vmap(rope, in_axes=(-2, -1), out_axes=(-2, -1)) + # RoPE.__call__ expects (..., seq_length, head_dim), so vmap over the heads + # axis and broadcast the positions, which are shared by every head. + apply = jax.vmap(rope, in_axes=(-2, None), out_axes=-2) query = apply(query, input_positions) key = apply(key, input_positions) return dot_product_attention(query, key, value, **kwargs) diff --git a/tests/nnx/nn/attention_test.py b/tests/nnx/nn/attention_test.py index ee581dfde..2527a1e6c 100644 --- a/tests/nnx/nn/attention_test.py +++ b/tests/nnx/nn/attention_test.py @@ -500,6 +500,41 @@ def apply_torch_rope(x): np.testing.assert_allclose(out_flax, out_torch, atol=1e-5) + def test_dot_product_attention_with_rope_input_positions(self): + """Explicit input_positions work and default positions match arange.""" + batch, seq_len, num_heads, head_dim = 2, 6, 4, 8 + + key = jax.random.key(0) + k1, k2, k3 = jax.random.split(key, 3) + query = jax.random.normal(k1, (batch, seq_len, num_heads, head_dim)) + kv = jax.random.normal(k2, (batch, seq_len, num_heads, head_dim)) + value = jax.random.normal(k3, (batch, seq_len, num_heads, head_dim)) + + rope = nnx.RoPE(theta=10000.0) + out_default = nnx.dot_product_attention_with_rope( + query, kv, value, rope=rope + ) + # Positions are shared by every head, so arange must match the default. + out_arange = nnx.dot_product_attention_with_rope( + query, kv, value, rope=rope, input_positions=jnp.arange(seq_len) + ) + np.testing.assert_allclose(out_arange, out_default, atol=1e-6) + + # RoPE encodes relative position: a uniform shift preserves every pairwise + # difference, so the attention output must not change. + out_shifted = nnx.dot_product_attention_with_rope( + query, kv, value, rope=rope, input_positions=jnp.arange(seq_len) + 3 + ) + np.testing.assert_allclose(out_shifted, out_default, atol=1e-5) + + # Non-uniform positions change the relative differences, so they must + # produce a different result — proving the positions take effect. + out_scaled = nnx.dot_product_attention_with_rope( + query, kv, value, rope=rope, input_positions=jnp.arange(seq_len) * 2 + ) + self.assertEqual(out_scaled.shape, out_default.shape) + self.assertFalse(jnp.allclose(out_scaled, out_default, atol=1e-3)) + def test_with_mha(self): """RoPE integrates correctly as attention_fn in MultiHeadAttention.""" with jax.numpy_rank_promotion("raise"):