zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-02 13:40:48 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

hard code freqs

Browse Source
This commit is contained in:
Awni Hannun
2024-08-22 16:48:46 -07:00
parent 0a52a9d55a
commit 9aabf08b23
1 changed files with 2 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
llms/mlx_lm/models/su_rope.py
View File

@@ -40,8 +40,7 @@ class SuScaledRotaryEmbedding(nn.Module):
"""
super().__init__()
freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
self._short_freqs = mx.array(short_factor, dtype=mx.float32) * freqs
self._long_freqs = mx.array(long_factor, dtype=mx.float32) * freqs
self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
self.original_max_position_embeddings = original_max_position_embeddings
self.scale = math.sqrt(
1
@@ -50,11 +49,6 @@ class SuScaledRotaryEmbedding(nn.Module):
)
def __call__(self, x, offset: int = 0):
freqs = (
self._long_freqs
if (offset + x.shape[2]) > self.original_max_position_embeddings
else self._short_freqs
)
return mx.fast.rope(
self.scale * x,
x.shape[-1],
@@ -62,5 +56,5 @@ class SuScaledRotaryEmbedding(nn.Module):
base=None,
scale=1.0,
offset=offset,
freqs=freqs,
freqs=self._freqs,
)