zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-21 12:06:42 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docs/src/dev/custom_metal_kernels.rst
View File

@ -128,6 +128,7 @@ relying on a copy from ``ensure_row_contiguous``:
input_names=["inp"], input_names=["inp"],
output_names=["out"], output_names=["out"],
source=source source=source
ensure_row_contiguous=False,
) )
def exp_elementwise(a: mx.array): def exp_elementwise(a: mx.array):
@ -138,7 +139,6 @@ relying on a copy from ``ensure_row_contiguous``:
threadgroup=(256, 1, 1), threadgroup=(256, 1, 1),
output_shapes=[a.shape], output_shapes=[a.shape],
output_dtypes=[a.dtype], output_dtypes=[a.dtype],
ensure_row_contiguous=False,
) )
return outputs[0] return outputs[0]

2
mlx/ops.cpp
View File

@ -2971,7 +2971,7 @@ array gather(
} }
for (auto& x : indices) { for (auto& x : indices) {
if (x.dtype() == bool_) { if (x.dtype() == bool_) {
throw("[Gather] Boolean indices not supported."); throw std::invalid_argument("[Gather] Boolean indices not supported.");
} }
} }

4
python/tests/test_fast.py
View File

@ -639,12 +639,12 @@ class TestFast(mlx_tests.MLXTestCase):
], ],
grid=(6, 1, 1), grid=(6, 1, 1),
threadgroup=(2, 1, 1), threadgroup=(2, 1, 1),
output_shapes=[(2, 2), (3, 2)], output_shapes=[(3, 2), (3, 2)],
output_dtypes=[mx.float32, mx.int32], output_dtypes=[mx.float32, mx.int32],
stream=mx.gpu, 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))) 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") @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")