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

Add more checks and clearer error messages to conv operations (#563)

Browse Source 
* Add more checks and clearer error messages to conv operations
This commit is contained in:
Jagrit Digani
2024-01-26 15:13:26 -08:00
committed by GitHub
parent 8fa6b322b9
commit bf17ab5002
3 changed files with 38 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -106,7 +106,9 @@ class RoPE(Module):
if (N, D, offset, base, scale, dtype) != cls._cos_sin_theta_key:
half_D = D // 2
positions = mx.arange(offset, N, dtype=dtype) * scale
freqs = mx.exp(-mx.arange(0.0, half_D, dtype=dtype) * (math.log(base) / half_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))