avoid mutating input in SuScaledRoPE and YarnRoPE (#1003)

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,

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()