From d4f4ff3c5ee2c33b7cf83c99065ccbb1c02e0eaf Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Thu, 25 Sep 2025 08:42:23 -0700
Subject: [PATCH] Allow None input to compiled functions (#2621)

* Allow None input to compiled functions

* Allow None input to compiled functions
---
 python/src/transforms.cpp    |  5 ++++-
 python/tests/test_compile.py | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp
index d88bc5f19..8c8c16316 100644
--- a/python/src/transforms.cpp
+++ b/python/src/transforms.cpp
@@ -441,6 +441,7 @@ struct PyCompiledFun {
     constexpr uint64_t array_identifier = 18446744073709551557UL;
     constexpr uint64_t list_identifier = 18446744073709551533UL;
     constexpr uint64_t dict_identifier = 18446744073709551521UL;
+    constexpr uint64_t none_identifier = 10239356951478402889UL;
 
     // Flatten the tree with hashed constants and structure
     std::function<void(nb::handle)> recurse;
@@ -476,10 +477,12 @@ struct PyCompiledFun {
       } else if (nb::isinstance<nb::float_>(obj)) {
         auto r = nb::cast<double>(obj);
         constants.push_back(*reinterpret_cast<uint64_t*>(&r));
+      } else if (obj.is_none()) {
+        constants.push_back(none_identifier);
       } else {
         std::ostringstream msg;
         msg << "[compile] Function arguments must be trees of arrays "
-            << "or constants (floats, ints, or strings), but received "
+            << "or constants (floats, ints, strings, or None), but received "
             << "type " << type_name_str(obj) << ".";
         throw std::invalid_argument(msg.str());
       }
diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py
index 5eb32ce4d..31fd38588 100644
--- a/python/tests/test_compile.py
+++ b/python/tests/test_compile.py
@@ -1037,6 +1037,20 @@ class TestCompile(mlx_tests.MLXTestCase):
         self.assertEqual(inner.__doc__, c_inner.__doc__)
         self.assertEqual(inspect.signature(inner), inspect.signature(c_inner))
 
+    def test_compile_with_none(self):
+        @mx.compile
+        def fun(x, y):
+            if y is None:
+                return mx.abs(x - 2.0)
+            else:
+                return mx.abs(x + y)
+
+        out = fun(mx.array(1.0), None)
+        self.assertEqual(out.item(), 1.0)
+
+        out = fun(mx.array(1.0), mx.array(2.0))
+        self.assertEqual(out.item(), 3.0)
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()