zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 08:20:15 +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());
}