zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-21 08:41:13 +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

4
python/src/transforms.cpp
View File

@ -477,13 +477,16 @@ auto py_compile(const py::function& fun) {
auto inputs = tree_flatten(args, true); auto inputs = tree_flatten(args, true);
// Get globally enclosed arrays so we don't compile through them // Get globally enclosed arrays so we don't compile through them
if (py::hasattr(fun, "__globals__")) {
auto global_inputs = tree_flatten(py::getattr(fun, "__globals__"), false); auto global_inputs = tree_flatten(py::getattr(fun, "__globals__"), false);
std::move( std::move(
std::begin(global_inputs), std::begin(global_inputs),
std::end(global_inputs), std::end(global_inputs),
std::back_inserter(inputs)); std::back_inserter(inputs));
}
// Get locally enclosed arrays so we don't compile through them // Get locally enclosed arrays so we don't compile through them
if (py::hasattr(fun, "__closure__")) {
auto closures = py::getattr(fun, "__closure__"); auto closures = py::getattr(fun, "__closure__");
if (py::isinstance<py::tuple>(closures)) { if (py::isinstance<py::tuple>(closures)) {
for (auto& closure : closures) { for (auto& closure : closures) {
@ -495,6 +498,7 @@ auto py_compile(const py::function& fun) {
std::back_inserter(inputs)); std::back_inserter(inputs));
} }
} }
}
// Compile and call // Compile and call
auto outputs = detail::compile(compile_fun, fun_id)(inputs); auto outputs = detail::compile(compile_fun, fun_id)(inputs);