From a1154ab94aced49ef01867fa6356cc3445e99d65 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 17 Mar 2026 00:43:33 -0700
Subject: [PATCH] Fix the rope mutation in a more natural way

---
 mlx_lm/models/rope_utils.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/mlx_lm/models/rope_utils.py b/mlx_lm/models/rope_utils.py
index 306f1f6..f3d5700 100644
--- a/mlx_lm/models/rope_utils.py
+++ b/mlx_lm/models/rope_utils.py
@@ -58,10 +58,8 @@ 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,
+            x.at[..., : self.dim].multiply(self._scale),
             self.dim,
             traditional=False,
             base=None,
@@ -183,8 +181,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]
+            x = x.at[..., : self.dims].multiply(self.mscale)
         return mx.fast.rope(
             x,
             self.dims,