enable and disable compiler

python/src/transforms.cpp

@@ -854,6 +854,24 @@ void init_transforms(py::module_& m) {
             function: A compiled function which has the same input arguments
             as ``fun`` and returns the the same output(s).
       )pbdoc");
+  m.def(
+      "disable_compiler",
+      &disable_compiler,
+      R"pbdoc(
+        disable_compiler() -> None
+
+        Globally disable compilation. Setting the environment variable
+        ``MLX_DISABLE_COMPILER`` can also be used to disable compilation.
+      )pbdoc");
+  m.def(
+      "enable_compiler",
+      &enable_compiler,
+      R"pbdoc(
+        enable_compiler() -> None
+
+        Globally enable compilation. This will override the environment
+        variable ``MLX_DISABLE_COMPILER`` if set.
+      )pbdoc");
 
   // Register static Python object cleanup before the interpreter exits
   auto atexit = py::module_::import("atexit");

python/tests/test_compile.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import io
 import unittest
 
 import mlx.core as mx
@@ -146,6 +147,32 @@ class TestCompile(mlx_tests.MLXTestCase):
         out = cfun(mx.array(3))
         self.assertEqual(out.item(), 4)
 
+    def test_enable_disable(self):
+        def fun(x):
+            y = x + 1
+            z = x + 1
+            return y + z
+
+        def count_prims(outputs):
+            buf = io.StringIO()
+            mx.export_to_dot(buf, outputs)
+            buf.seek(0)
+            return len([l for l in buf.read().split() if "label" in l])
+
+        x = mx.array(1.0)
+        cfun = mx.compile(fun)
+        n_compiled = count_prims(cfun(x))
+
+        # Check disabled
+        mx.disable_compiler()
+        n_uncompiled = count_prims(cfun(x))
+        self.assertTrue(n_compiled < n_uncompiled)
+
+        # Check renabled
+        mx.enable_compiler()
+        n_enable_compiled = count_prims(cfun(x))
+        self.assertEqual(n_compiled, n_enable_compiled)
+
 
 if __name__ == "__main__":
     unittest.main()