zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 01:17:26 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix test

Browse Source
This commit is contained in:
Alex Barron 2024-12-06 10:26:54 -08:00
parent 769704653a
commit 82a956c1d9
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
python/tests/test_fast_sdpa.py
View File

@ -128,10 +128,13 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
B = 1 B = 1
H = 32 H = 32
dtypes = [mx.float32]
if self.is_apple_silcon:
dtypes.append(mx.float16)
tests = product( tests = product(
[1, 7, 9, 32, 63, 67, 129, 2000], # sequence length [1, 7, 9, 32, 63, 67, 129, 2000], # sequence length
[False, True], # gqa [False, True], # gqa
[mx.float32, mx.float16], dtypes,
[4, 8], # bits [4, 8], # bits
) )
for sequence_length, do_gqa, dtype, bits in tests: for sequence_length, do_gqa, dtype, bits in tests: