throw for certain cases of non captured inputs in compile (#1401)

python/src/fast.cpp

@@ -302,20 +302,20 @@ void init_fast(nb::module_& parent_module) {
       A jit-compiled custom Metal kernel defined from a source string.
 
       Args:
-       name (str): Name for the kernel.
-       input_names (List[str]): The parameter names of the inputs in the
-          function signature.
-       output_names (List[str]): The parameter names of the outputs in the
+        name (str): Name for the kernel.
+        input_names (List[str]): The parameter names of the inputs in the
            function signature.
-       source (str): Source code. This is the body of a function in Metal,
-           the function signature will be automatically generated.
-       header (str): Header source code to include before the main function.
-           Useful for helper functions or includes that should live outside of
-            the main function body.
-       ensure_row_contiguous (bool): Whether to ensure the inputs are row contiguous
-           before the kernel runs. Default: ``True``.
-       atomic_outputs (bool): Whether to use atomic outputs in the function signature
-           e.g. ``device atomic<float>``. Default: ``False``.
+        output_names (List[str]): The parameter names of the outputs in the
+            function signature.
+        source (str): Source code. This is the body of a function in Metal,
+            the function signature will be automatically generated.
+        header (str): Header source code to include before the main function.
+            Useful for helper functions or includes that should live outside of
+             the main function body.
+        ensure_row_contiguous (bool): Whether to ensure the inputs are row contiguous
+            before the kernel runs. Default: ``True``.
+        atomic_outputs (bool): Whether to use atomic outputs in the function signature
+            e.g. ``device atomic<float>``. Default: ``False``.
 
       Returns:
         Callable ``metal_kernel``.

python/tests/test_compile.py

@@ -733,6 +733,31 @@ class TestCompile(mlx_tests.MLXTestCase):
             expected = fn(x)
             self.assertTrue(mx.array_equal(expected, out))
 
+    def test_compile_without_captured_inputs(self):
+        x = mx.array([1, 2, 3]) + 2
+
+        def fn(a):
+            y = x + 1
+            return a + y
+
+        with self.assertRaises(ValueError):
+            y = mx.compile(fn)(x)
+
+        x = mx.array([1.0, 2.0]) + mx.array([1.0, 2.0])
+        y = None
+
+        def fn(x):
+            nonlocal y
+            if y is None:
+                y = mx.array([1.0, 2.0])
+
+            y = y + x
+            return y
+
+        fn(x)
+        with self.assertRaises(ValueError):
+            y = mx.compile(fn)(x)
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_fast.py

@@ -689,6 +689,27 @@ class TestFast(mlx_tests.MLXTestCase):
         )
         self.assertTrue(mx.allclose(out[0], mx.exp(a)))
 
+    @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
+    def test_custom_kernel_attributes(self):
+        a = mx.zeros(shape=(1, 1))
+        kernel = mx.fast.metal_kernel(
+            name="test_fun",
+            input_names=["a"],
+            output_names=["out"],
+            source="""
+                out[0] = threads_per_threadgroup.x;
+            """,
+        )
+        out = kernel(
+            inputs=[a],
+            grid=(2, 1, 1),
+            threadgroup=(2, 1, 1),
+            output_shapes=[(1, 1)],
+            output_dtypes=[mx.uint32],
+            stream=mx.gpu,
+        )[0]
+        self.assertEqual(out.item(), 2)
+
 
 if __name__ == "__main__":
     unittest.main()