Rename Qwen3NextRotaryEmbedding to PartialRotaryEmbedding to make more generic and reusable.

Author: gagika

src/maxtext/layers/attentions.py

@@ -62,7 +62,7 @@
     Qwen3OmniMoeVisionRotaryEmbedding,
     RotaryEmbedding,
     YarnRotaryEmbedding,
-    Qwen3NextRotaryEmbedding,
+    PartialRotaryEmbedding,
 )
 from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned, default_bias_init
 from maxtext.layers.linears import DenseGeneral, canonicalize_tuple, normalize_axes
@@ -814,7 +814,7 @@ def init_rotary_embedding(self):
           rngs=self.rngs,
       )
     elif self.is_qwen3_next:
-      rotary_embedding = Qwen3NextRotaryEmbedding(
+      rotary_embedding = PartialRotaryEmbedding(
           min_timescale=self.config.rope_min_timescale,
           max_timescale=self.config.rope_max_timescale,
           mesh=self.mesh,

src/maxtext/layers/embeddings.py

@@ -406,7 +406,7 @@ def llama_rotary_embedding_as_linen(
   )
 
 
-def qwen3_next_rotary_embedding_as_linen(
+def partial_rotary_embedding_as_linen(
     *,
     min_timescale: int,
     max_timescale: int,
@@ -418,7 +418,7 @@ def qwen3_next_rotary_embedding_as_linen(
     shard_mode: ShardMode = ShardMode.AUTO,
     name: str | None = None,
 ):
-  """Initializes the Qwen3NextRotaryEmbedding module and returns it as a Linen module.
+  """Initializes the PartialRotaryEmbedding module and returns it as a Linen module.
 
   Args:
     min_timescale: Start of the geometric index. Determines the periodicity of
@@ -432,7 +432,7 @@ def qwen3_next_rotary_embedding_as_linen(
     name: Name of the Linen module.
   """
   return nnx_wrappers.to_linen(
-      Qwen3NextRotaryEmbedding,
+      PartialRotaryEmbedding,
       min_timescale=min_timescale,
       max_timescale=max_timescale,
       mesh=mesh,
@@ -446,8 +446,8 @@ def qwen3_next_rotary_embedding_as_linen(
   )
 
 
-class Qwen3NextRotaryEmbedding(RotaryEmbedding):
-  """Qwen3 Next variant of ROPE (partial ROPE)"""
+class PartialRotaryEmbedding(RotaryEmbedding):
+  """Rotary Position Embedding applied to a partial fraction of dimensions."""
 
   def __init__(
       self,
@@ -461,7 +461,7 @@ def __init__(
       shard_mode: ShardMode = ShardMode.AUTO,
       rngs: nnx.Rngs = None,
   ):
-    """Initializes the Qwen3NextRotaryEmbedding module.
+    """Initializes the PartialRotaryEmbedding module.
 
     Args:
       min_timescale: Start of the geometric index. Determines the periodicity of
@@ -476,6 +476,7 @@ def __init__(
     self.partial_rotary_factor = partial_rotary_factor
     self.rotary_dim = int(self.head_dim * self.partial_rotary_factor)
 
+    # Initialize the base class with only the rotary_dim
     super().__init__(
         min_timescale=min_timescale,
         max_timescale=max_timescale,
@@ -488,7 +489,7 @@ def __init__(
     )
 
   def __call__(self, inputs: jax.Array, position: None | jax.Array = None) -> jax.Array:
-    """Applies LLaMA variant of rotary position embedding.
+    """Applies Partial variant of rotary position embedding.
 
     Args:
       inputs: The input sequence on which to apply the Rotary position
@@ -499,6 +500,7 @@ def __call__(self, inputs: jax.Array, position: None | jax.Array = None) -> jax.
     Returns:
       A jax.Array of shape [B, S, H, D - rotary_dim] with rotary position embeddings applied.
     """
+    # Split, apply base RoPE to the first fraction, and concatenate
     inputs_rot, inputs_pass = jnp.split(inputs, [self.rotary_dim], axis=-1)
     inputs_rot = super().__call__(inputs_rot, position)
     inputs = jnp.concatenate([inputs_rot, inputs_pass], axis=-1)

src/maxtext/models/qwen3.py

@@ -615,7 +615,7 @@ class Qwen3NextFullAttention(nnx.Module):
     - Query and Gate splitting from a single q projection.
     - Application of a sigmoid gate to the attention output.
     - Usage of `Qwen3NextRMSNorm` for query and key normalization.
-    - Usage of `Qwen3NextRotaryEmbedding` for partial rotary position embeddings.
+    - Usage of `PartialRotaryEmbedding` for partial rotary position embeddings.
       - Partial ROPE is applied to the first 25% of head dimensions
 
   Attributes:

tests/unit/partial_rotary_embedding_test.py

@@ -0,0 +1,166 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the partial rotary position embedding layer.
+
+The new PartialRotaryEmbedding class is a thin wrapper around
+`RotaryEmbedding` that applies RoPE only to the first fraction of the
+hidden dimensions.  The tests below exercise the half/fully-rotated
+cases and verify basic shift invariance in the same style used by our
+existing rotary unit tests.
+"""
+
+import sys
+import unittest
+
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from flax import nnx
+import numpy as np
+
+from maxtext.layers.embeddings import PartialRotaryEmbedding, RotaryEmbedding
+from maxtext.configs import pyconfig
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+
+class PartialRotaryEmbeddingTest(unittest.TestCase):
+  """Tests for the PartialRotaryEmbedding layer."""
+
+  def setUp(self):
+    super().setUp()
+    # build a simple config and mesh like other embedding tests
+    self.cfg = pyconfig.initialize(
+        [sys.argv[0], get_test_config_path()],
+        run_name="test_embeddings",
+        enable_checkpointing=False,
+    )
+    devices_array = maxtext_utils.create_device_mesh(self.cfg)
+    self.mesh = Mesh(devices_array, self.cfg.mesh_axes)
+    self.nnx_rng = nnx.Rngs(params=0)
+
+  def test_partial_rotary_half(self):
+    """The first half of the hidden dim should be rotated, the rest passthrough."""
+    batch_size, seq_len, num_heads, head_dim = 2, 16, 4, 8
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, seq_len, num_heads, head_dim))
+    positions = jnp.arange(seq_len, dtype=jnp.int32).reshape(1, seq_len)
+
+    rope_half = PartialRotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=head_dim,
+        partial_rotary_factor=0.5,
+        rngs=self.nnx_rng,
+        cast_as_fprop_dtype=False,
+    )
+
+    y_half = rope_half(inputs, positions)
+
+    rotary_dim = head_dim // 2
+    inputs_rot, inputs_pass = inputs[..., :rotary_dim], inputs[..., rotary_dim:]
+
+    rope_full_for_rot_part = RotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=rotary_dim,
+        rngs=self.nnx_rng,
+        cast_as_fprop_dtype=False,
+    )
+    y_rot_expected = rope_full_for_rot_part(inputs_rot, positions)
+
+    np.testing.assert_allclose(
+        y_half[..., :rotary_dim],
+        y_rot_expected,
+        rtol=1e-6,
+        atol=1e-6,
+        err_msg="First fraction should be rotated.",
+    )
+    np.testing.assert_allclose(
+        y_half[..., rotary_dim:],
+        inputs_pass,
+        rtol=1e-6,
+        atol=1e-6,
+        err_msg="Remaining dims should pass through.",
+    )
+
+  def test_partial_rotary_full(self):
+    """A partial factor of 1.0 is equivalent to the base rotary embedding."""
+    batch_size, seq_len, num_heads, head_dim = 1, 4, 4, 8
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, seq_len, num_heads, head_dim))
+    positions = jnp.arange(seq_len, dtype=jnp.int32).reshape(1, seq_len)
+
+    rope_partial = PartialRotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=head_dim,
+        partial_rotary_factor=1.0,
+        rngs=self.nnx_rng,
+    )
+    y_partial = rope_partial(inputs, positions)
+
+    rope_full = RotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=head_dim,
+        rngs=self.nnx_rng,
+    )
+    y_full = rope_full(inputs, positions)
+
+    np.testing.assert_allclose(
+        y_partial,
+        y_full,
+        rtol=1e-6,
+        atol=1e-6,
+        err_msg="PartialRotaryEmbedding with factor=1 should equal full rotary.",
+    )
+
+  def test_shift_invariance(self):
+    """Verify that rotary attention computed from partial embedding is shift invariant."""
+    batch_size, seq_len, num_heads, head_dim = 1, 20, 4, 8
+    inputs = jax.random.normal(jax.random.PRNGKey(0), (batch_size, seq_len, num_heads, head_dim))
+    positions = jnp.arange(seq_len, dtype=jnp.int32).reshape(1, seq_len)
+
+    rope = PartialRotaryEmbedding(
+        min_timescale=1,
+        max_timescale=10000,
+        mesh=self.mesh,
+        embedding_dims=head_dim,
+        partial_rotary_factor=0.5,
+        rngs=self.nnx_rng,
+        cast_as_fprop_dtype=False,
+    )
+
+    def get_attn(pos):
+      y = rope(inputs, pos)
+      return np.einsum("BSNH,BTNH->BSTN", y, y)
+
+    ref_attn = get_attn(positions)
+    shifted_attn = get_attn(positions + 3)
+
+    np.testing.assert_allclose(
+        ref_attn,
+        shifted_attn,
+        rtol=1e-6,
+        atol=1e-6,
+        err_msg="PartialRotaryEmbedding attention should be shift-invariant.",
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()