diff --git a/python/src/array.cpp b/python/src/array.cpp index afb691428..9115ada6e 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -39,6 +39,10 @@ py::list to_list(array& a, size_t index, int dim) { } auto to_scalar(array& a) { + { + py::gil_scoped_release nogil; + a.eval(); + } switch (a.dtype()) { case bool_: return py::cast(a.item()); @@ -73,7 +77,10 @@ py::object tolist(array& a) { if (a.ndim() == 0) { return to_scalar(a); } - a.eval(); + { + py::gil_scoped_release nogil; + a.eval(); + } py::object pl; switch (a.dtype()) { case bool_: @@ -644,6 +651,7 @@ void init_array(py::module_& m) { .def_buffer([](array& a) { // Eval if not already evaled if (!a.is_evaled()) { + py::gil_scoped_release nogil; a.eval(); } return pybind11::buffer_info( @@ -942,6 +950,7 @@ void init_array(py::module_& m) { "__repr__", [](array& a) { if (!a.is_evaled()) { + py::gil_scoped_release nogil; a.eval(); } std::ostringstream os; diff --git a/python/src/load.cpp b/python/src/load.cpp index 92ad5e808..18e89c7fb 100644 --- a/python/src/load.cpp +++ b/python/src/load.cpp @@ -195,6 +195,8 @@ mlx_load_gguf_helper(py::object file, StreamOrDevice s) { std::unordered_map mlx_load_npz_helper( py::object file, StreamOrDevice s) { + bool own_file = py::isinstance(file); + py::module_ zipfile = py::module_::import("zipfile"); if (!is_zip_file(zipfile, file)) { throw std::invalid_argument( @@ -223,9 +225,11 @@ std::unordered_map mlx_load_npz_helper( } // If we don't own the stream and it was passed to us, eval immediately - for (auto& [key, arr] : array_dict) { + if (!own_file) { py::gil_scoped_release gil; - arr.eval(); + for (auto& [key, arr] : array_dict) { + arr.eval(); + } } return array_dict; @@ -260,7 +264,7 @@ LoadOutputTypes mlx_load_helper( fname = file.attr("name").cast(); } else { throw std::invalid_argument( - "[load] Input must be a file-like object, or string"); + "[load] Input must be a file-like object opened in binary mode, or string"); } size_t ext = fname.find_last_of('.'); if (ext == std::string::npos) { @@ -432,7 +436,7 @@ void mlx_savez_helper( auto py_ostream = zipfile_object.open(fname, 'w'); auto writer = std::make_shared(py_ostream); { - py::gil_scoped_release gil; + py::gil_scoped_release nogil; save(writer, a); } } @@ -443,20 +447,20 @@ void mlx_savez_helper( void mlx_save_safetensor_helper(py::object file, py::dict d) { auto arrays_map = d.cast>(); if (py::isinstance(file)) { - save_safetensors(py::cast(file), arrays_map); - return; + { + py::gil_scoped_release nogil; + save_safetensors(py::cast(file), arrays_map); + } } else if (is_ostream_object(file)) { auto writer = std::make_shared(file); { - py::gil_scoped_release gil; + py::gil_scoped_release nogil; save_safetensors(writer, arrays_map); } - - return; + } else { + throw std::invalid_argument( + "[save_safetensors] Input must be a file-like object, or string"); } - - throw std::invalid_argument( - "[save_safetensors] Input must be a file-like object, or string"); } void mlx_save_gguf_helper( @@ -468,12 +472,17 @@ void mlx_save_gguf_helper( if (m) { auto metadata_map = m.value().cast>(); - save_gguf(py::cast(file), arrays_map, metadata_map); + { + py::gil_scoped_release nogil; + save_gguf(py::cast(file), arrays_map, metadata_map); + } } else { - save_gguf(py::cast(file), arrays_map); + { + py::gil_scoped_release nogil; + save_gguf(py::cast(file), arrays_map); + } } - return; + } else { + throw std::invalid_argument("[save_gguf] Input must be a string"); } - - throw std::invalid_argument("[save_safetensors] Input must be a string"); } diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index 8706706d6..0c15daf85 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -509,7 +509,10 @@ void init_transforms(py::module_& m) { "eval", [](const py::args& args) { std::vector arrays = tree_flatten(args); - eval(arrays); + { + py::gil_scoped_release nogil; + eval(arrays); + } }, R"pbdoc( eval(*args) -> None