zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-26 15:58:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Allow None input to compiled functions (#2621)

Browse Source 
* Allow None input to compiled functions

* Allow None input to compiled functions
This commit is contained in:
Awni Hannun
2025-09-25 08:42:23 -07:00
committed by GitHub
parent 7c7e48dbd1
commit d4f4ff3c5e
2 changed files with 18 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/src/transforms.cpp
View File

@@ -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());
}

14
python/tests/test_compile.py
View File

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