zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-19 23:51:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix test

Browse Source
This commit is contained in:
Awni Hannun 2024-01-15 11:30:24 -08:00
parent 4f50935c2c
commit 6189111494
1 changed files with 18 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
python/src/transforms.cpp
View File

@ -477,22 +477,26 @@ auto py_compile(const py::function& fun) {
auto inputs = tree_flatten(args, true);
// Get globally enclosed arrays so we don't compile through them
auto global_inputs = tree_flatten(py::getattr(fun, "__globals__"), false);
std::move(
std::begin(global_inputs),
std::end(global_inputs),
std::back_inserter(inputs));
if (py::hasattr(fun, "__globals__")) {
auto global_inputs = tree_flatten(py::getattr(fun, "__globals__"), false);
std::move(
std::begin(global_inputs),
std::end(global_inputs),
std::back_inserter(inputs));
}
// Get locally enclosed arrays so we don't compile through them
auto closures = py::getattr(fun, "__closure__");
if (py::isinstance<py::tuple>(closures)) {
for (auto& closure : closures) {
auto enclosed_inputs =
tree_flatten(py::getattr(closure, "cell_contents"), false);
std::move(
std::begin(enclosed_inputs),
std::end(enclosed_inputs),
std::back_inserter(inputs));
if (py::hasattr(fun, "__closure__")) {
auto closures = py::getattr(fun, "__closure__");
if (py::isinstance<py::tuple>(closures)) {
for (auto& closure : closures) {
auto enclosed_inputs =
tree_flatten(py::getattr(closure, "cell_contents"), false);
std::move(
std::begin(enclosed_inputs),
std::end(enclosed_inputs),
std::back_inserter(inputs));
}
}
}