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

fix custom kernel test (#2510)

Browse Source
This commit is contained in:
Awni Hannun
2025-08-18 06:45:59 -07:00
committed by GitHub
parent 1df9887998
commit c5fcd5b61b
3 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/tests/test_fast.py
View File

@@ -639,12 +639,12 @@ class TestFast(mlx_tests.MLXTestCase):
],
grid=(6, 1, 1),
threadgroup=(2, 1, 1),
output_shapes=[(2, 2), (3, 2)],
output_shapes=[(3, 2), (3, 2)],
output_dtypes=[mx.float32, mx.int32],
stream=mx.gpu,
)
self.assertTrue(mx.allclose(out[0], mx.full((2, 2), 14.0484)))
self.assertTrue(mx.allclose(out[0], mx.full((3, 2), 14.0484)))
self.assertTrue(mx.allclose(out[1], mx.full((3, 2), -2, dtype=mx.int32)))
@unittest.skipIf(not mx.metal.is_available(), "Metal is not available")