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

Added implementation for Scaled RoPE. (#261)

Browse Source 
* Added scale for RoPE

* Ran pre-commit

* Added RoPE scaling test

* Added docstring for scale parameter

* Modified docstrings
This commit is contained in:
Hazem Essam
2023-12-31 16:06:01 +02:00
committed by GitHub
parent a020a2d49d
commit e3b8da2a49
2 changed files with 18 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/tests/test_nn.py
View File

@@ -749,7 +749,7 @@ class TestNN(mlx_tests.MLXTestCase):
self.assertEqual(y.dtype, mx.float32)
def test_rope(self):
for kwargs in [{}, {"traditional": False}, {"base": 10000}]:
for kwargs in [{}, {"traditional": False}, {"base": 10000}, {"scale": 0.25}]:
rope = nn.RoPE(4, **kwargs)
shape = (1, 3, 4)
x = mx.random.uniform(shape=shape)