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

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