feat(nn): add freqs and partial_rotary_factor support to RoPE

[shadyuwugurl]

Summary

mx.fast.rope already accepts a caller-supplied freqs array, but mlx.nn.RoPE did not expose it — it only forwarded dims, traditional, base and scale. As a result, every rotary encoding that is not a plain power law had to be reimplemented in Python on top of MLX.

This PR exposes freqs on nn.RoPE and adds a partial_rotary_factor convenience that builds the wavelengths for a proportional (partial) rotary encoding.

rope = nn.RoPE(256, base=10000.0, partial_rotary_factor=0.25)   # partial rotary
rope = nn.RoPE(256, freqs=my_wavelengths)                        # arbitrary scaling

Why this matters

mlx-lm currently carries six Python reimplementations of rotary encodings in mlx_lm/models/rope_utils.py — SuScaledRoPE, Llama3RoPE, YarnRoPE, ProportionalRoPE, DynamicNTKScalingRoPE — purely because nn.RoPE could not express them. All of them are thin wrappers that build a freqs array and call mx.fast.rope.

Current models that need this:

Model	Rotary variant	Needed
google/gemma-4-12B (gemma4_text)	partial_rotary_factor: 0.25, rope_type: "proportional", global_head_dim: 512 vs head_dim: 256	partial + custom dims
Qwen3.5 (qwen3_5)	partial_rotary_factor: 0.25, mrope_section, linear_key_head_dim: 192	partial

Semantics

partial_rotary_factor follows transformers' _compute_proportional_rope_parameters exactly:

rotated_dims = int(dims * partial_rotary_factor)
exponents    = arange(0, rotated_dims, 2) / dims          # denominator is the FULL dims
freqs        = concat([base ** exponents, inf * ((dims - rotated_dims) // 2)])

The non-rotated features get an infinite wavelength, so their inverse frequency is zero and mx.fast.rope leaves them unchanged. This matches the HF reference, where inv_freq is zero-padded and the cosine/sine pair is duplicated.

The frequency-scaling factor that HF applies via rope_scaling["factor"] is expressed by the freqs path (freqs *= factor), keeping the two arguments independent.

Behaviour

• The default path is unchanged — nn.RoPE(dims) produces exactly the same result as before.
• freqs takes precedence over base (mirroring mx.fast.rope, which requires exactly one of the two).
• freqs is stored with object.__setattr__ so it is not registered as a module parameter; nn.RoPE(8, freqs=...).parameters() is empty, as it should be for a constant.

Validation

import mlx.core as mx
from mlx.nn import RoPE

# custom frequencies match the fast op exactly
freqs = mx.array([10000.0, 1000.0, 100.0, 10.0])
x = mx.random.uniform(shape=(2, 4, 6, 8))
assert mx.allclose(RoPE(8, freqs=freqs)(x),
                   mx.fast.rope(x, 8, traditional=False, base=None, scale=1.0, offset=0, freqs=freqs))

# partial rotary matches a NumPy reference of the HF convention
# (split-half pairing, zero inverse frequency on the untouched features)

New tests in python/tests/test_nn.py (TestRoPEVariants) cover:

• default path still matches mx.fast.rope
• custom freqs match mx.fast.rope
• partial rotary matches an independent NumPy reference
• partial_rotary_factor and explicit freqs agree
• non-rotated features are bit-identical to the input
• offsets and dtype preservation
• freqs are not exposed as parameters
• invalid arguments raise ValueError

Local run on an M5 (Apple Silicon), mlx 0.32.4:

$ python -m pytest python/tests/test_nn.py -q
81 passed, 1 skipped, 3 subtests passed

Note on MLX_ENABLE_TF32

While validating numerical parity of the rotary implementation I found that float32 matmuls on the Metal GPU path use TF32 by default on M5 (mx.utils.h, enable_tf32(), default 1), giving ~7e-4 relative error versus ~1e-6 for the same op on MLX's CPU stream. With MLX_ENABLE_TF32=0 the error drops to ~5e-7. This is orthogonal to this PR and is already discussed in #4595 — I mention it only because it changes the expected tolerance of float32 comparisons on Apple Silicon.

🤖 Generated with Claude Code