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")