Compile float64 functions on CPU (#2311)

python/tests/test_compile.py

@@ -2,6 +2,7 @@
 
 import gc
 import io
+import math
 import unittest
 from functools import partial
 
@@ -979,6 +980,17 @@ class TestCompile(mlx_tests.MLXTestCase):
 
         self.assertEqual(mem_pre, mem_post)
 
+    def test_double_constant(self):
+        with mx.stream(mx.cpu):
+            x = mx.array(1.0, dtype=mx.float64)
+
+            def fun(x):
+                return (x + math.pi) * 2.0
+
+            y = fun(x).item()
+            y_compiled = mx.compile(fun)(x).item()
+            self.assertEqual(y, y_compiled)
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()