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