From c5fcd5b61b2483a157388dae12099e6c71c33713 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Mon, 18 Aug 2025 06:45:59 -0700 Subject: [PATCH] fix custom kernel test (#2510) --- docs/src/dev/custom_metal_kernels.rst | 2 +- mlx/ops.cpp | 2 +- python/tests/test_fast.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/src/dev/custom_metal_kernels.rst b/docs/src/dev/custom_metal_kernels.rst index 873b1e544..1febe960a 100644 --- a/docs/src/dev/custom_metal_kernels.rst +++ b/docs/src/dev/custom_metal_kernels.rst @@ -128,6 +128,7 @@ relying on a copy from ``ensure_row_contiguous``: input_names=["inp"], output_names=["out"], source=source + ensure_row_contiguous=False, ) def exp_elementwise(a: mx.array): @@ -138,7 +139,6 @@ relying on a copy from ``ensure_row_contiguous``: threadgroup=(256, 1, 1), output_shapes=[a.shape], output_dtypes=[a.dtype], - ensure_row_contiguous=False, ) return outputs[0] diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 14e135deb..c8583c72f 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2971,7 +2971,7 @@ array gather( } for (auto& x : indices) { if (x.dtype() == bool_) { - throw("[Gather] Boolean indices not supported."); + throw std::invalid_argument("[Gather] Boolean indices not supported."); } } diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 13c65de99..f79a62a15 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -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")