Fix the rope mutation in a more natural way

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,