From 2146e4ed184726082addc26992cd5e2b107efe9d Mon Sep 17 00:00:00 2001
From: mm65x <martino.mansoldo@gmail.com>
Date: Mon, 16 Mar 2026 01:13:47 +0000
Subject: [PATCH] avoid mutating input in SuScaledRoPE and YarnRoPE (#1003)

---
 mlx_lm/models/rope_utils.py |  3 ++-
 tests/test_models.py        | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/models/rope_utils.py b/mlx_lm/models/rope_utils.py
index a729a50..306f1f6 100644
--- a/mlx_lm/models/rope_utils.py
+++ b/mlx_lm/models/rope_utils.py
@@ -58,6 +58,7 @@ class SuScaledRoPE(nn.Module):
         self._scale = long_mscale or (1.0 if factor <= 1.0 else default_scale(factor))
 
     def __call__(self, x, offset: Union[int, mx.array] = 0):
+        x = x[...]
         x[..., : self.dim] = self._scale * x[..., : self.dim]
         return mx.fast.rope(
             x,
@@ -71,7 +72,6 @@ class SuScaledRoPE(nn.Module):
 
 
 class Llama3RoPE(nn.Module):
-
     def __init__(
         self,
         dims: int,
@@ -183,6 +183,7 @@ class YarnRoPE(nn.Module):
 
     def __call__(self, x, offset=0):
         if self.mscale != 1.0:
+            x = x[...]
             x[..., : self.dims] = self.mscale * x[..., : self.dims]
         return mx.fast.rope(
             x,
diff --git a/tests/test_models.py b/tests/test_models.py
index bcc63bc..0b2a963 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -242,6 +242,30 @@ class TestModels(unittest.TestCase):
         )
         self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
 
+    def test_su_scaled_rope_no_mutation(self):
+        rope = rope_utils.SuScaledRoPE(
+            dims=8,
+            max_position_embeddings=131072,
+            original_max_position_embeddings=4096,
+            long_factor=[1.0] * 4,
+        )
+        x = mx.ones((1, 2, 4, 8))
+        rope(x)
+        mx.eval(x)
+        self.assertTrue((x == 1).all())
+
+    def test_yarn_rope_no_mutation(self):
+        rope = rope_utils.YarnRoPE(
+            dims=8,
+            scaling_factor=2.0,
+            mscale=1.0,
+            mscale_all_dim=0,
+        )
+        x = mx.ones((1, 2, 4, 8))
+        rope(x)
+        mx.eval(x)
+        self.assertTrue((x == 1).all())
+
     def test_quantized_sdpa(self):
         cache = KVCache()