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

fix test + choose stream with slight care

Browse Source
This commit is contained in:
Awni Hannun 2024-01-13 13:34:27 -08:00
parent 9051fa1eaa
commit fd94be28ea
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
mlx/transforms.cpp
View File

@ -204,11 +204,17 @@ void eval(const std::vector<array>& outputs) {
std::unordered_set<std::uintptr_t> cache; std::unordered_set<std::uintptr_t> cache;
std::unordered_map<std::uintptr_t, std::shared_future<void>> deps; std::unordered_map<std::uintptr_t, std::shared_future<void>> deps;
auto synchronizer = array( // Make an effort to choose a good output stream
{}, Stream stream = default_stream(default_device());
bool_, for (auto& o : outputs) {
std::make_unique<Synchronizer>(default_stream(default_device())), if (!o.is_evaled() && o.has_primitive()) {
outputs); stream = o.primitive().stream();
break;
}
}
auto synchronizer =
array({}, bool_, std::make_unique<Synchronizer>(stream), outputs);
recurse = [&](const array& a) { recurse = [&](const array& a) {
auto id = a.id(); auto id = a.id();