From e64349bbdd2d87781edf06c9a03d75adf334c746 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 27 Aug 2024 17:01:22 -0700
Subject: [PATCH] Make eval just wait if all arrays are scheduled (#1368)

---
 mlx/transforms.cpp | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp
index f8d5648f8..19387388e 100644
--- a/mlx/transforms.cpp
+++ b/mlx/transforms.cpp
@@ -186,10 +186,35 @@ array eval_impl(std::vector<array> outputs, bool async) {
 }
 
 void async_eval(std::vector<array> outputs) {
+  if (outputs.empty()) {
+    return;
+  }
+
+  if (std::none_of(outputs.begin(), outputs.end(), [](array& x) {
+        return x.status() == array::Status::unscheduled;
+      })) {
+    return;
+  }
+
   eval_impl(std::move(outputs), true);
 }
 
 void eval(std::vector<array> outputs) {
+  if (outputs.empty()) {
+    return;
+  }
+
+  if (std::none_of(outputs.begin(), outputs.end(), [](array& x) {
+        return x.status() == array::Status::unscheduled;
+      })) {
+    for (auto& x : outputs) {
+      if (!x.is_available()) {
+        x.event().wait();
+      }
+    }
+    return;
+  }
+
   eval_impl(std::move(outputs), false).event().wait();
 }