Fix the rope mutation in a more natural way

2026-03-18 13:34:13 -07:00
1 changed files with 2 additions and 5 deletions
@@ -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,