zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-04 08:11:13 +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_map<std::uintptr_t, std::shared_future<void>> deps;
auto synchronizer = array(
{},
bool_,
std::make_unique<Synchronizer>(default_stream(default_device())),
outputs);
// Make an effort to choose a good output stream
Stream stream = default_stream(default_device());
for (auto& o : outputs) {
if (!o.is_evaled() && o.has_primitive()) {
stream = o.primitive().stream();
break;
}
}
auto synchronizer =
array({}, bool_, std::make_unique<Synchronizer>(stream), outputs);
recurse = [&](const array& a) {
auto id = a.id();