zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-23 01:51:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix cache key in RoPE (#561)

Browse Source
This commit is contained in:
David Koski 2024-01-26 13:10:02 -08:00 committed by GitHub
parent 077c1ee64a
commit 874b739f3c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
python/mlx/nn/layers/positional_encoding.py
View File

@ -104,9 +104,9 @@ class RoPE(Module):
dtype=mx.float32,
):
if (N, D, offset, base, scale, dtype) != cls._cos_sin_theta_key:
D = D // 2
half_D = D // 2
positions = mx.arange(offset, N, dtype=dtype) * scale
freqs = mx.exp(-mx.arange(0.0, D, dtype=dtype) * (math.log(base) / D))
freqs = mx.exp(-mx.arange(0.0, half_D, dtype=dtype) * (math.log(base) / half_D))
theta = mx.reshape(positions, (-1, 1)) * mx.reshape(freqs, (1, -1))
cls._cos_sin_theta_key = (N, D, offset, base, scale, dtype)
cls._cos_sin_theta_value = (mx.cos(theta), mx.sin(theta))