Repository: ml-explore/mlx (2026-06-06T04:07:59Z)

Issue:
compile: failed first trace leaves a stale cache entry that returns empty-success on retry
### Describe the bug

The caching lambda built by `mlx::core::detail::compile` marks a cache entry **non-empty before it traces it**. In `mlx/compile.cpp`, the per-call lambda does (paraphrased):

```cpp
auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants);
if (entry.empty) {
  entry.empty = false;                 // <-- marked filled BEFORE the trace
  entry.constants = std::move(constants);
  std::tie(entry.inputs, entry.outputs, entry.extra) =
      compile_trace(fun, inputs, shapeless);   // <-- may throw
  ...
}
```

If `compile_trace` throws on the **first** trace — e.g. the traced `fun` raises — the entry is left `empty == false` but with no `inputs` / `outputs` / tape filled. A later call with matching inputs then `find()`s that entry, sees `empty == false`, **skips tracing**, and `compile_replace`s an effectively empty tape — returning empty outputs as a spurious success instead of re-raising (or re-tracing).

### To reproduce

Compile a function that raises on its first invocation, then call it twice with the same (matching) inputs:

1. First call: `compile_trace` raises and propagates (expected).
2. Second call: the half-filled entry is hit, the trace is skipped, and **empty outputs are returned as a success**.

A nullary compiled function is the sharpest case (its inputs always match, so the second call is a guaranteed `find` on the same entry).

### Expected behavior

A failed first trace should not leave a usable-looking cache entry. Either:

- set `entry.empty = false` only **after** `compile_trace` (and the dfs / simplify / fuse) complete successfully, or
- roll the entry back (`entry.empty = true`, or erase it) if tracing throws,

so a retry re-traces cleanly and a failure stays a failure.

### Why it matters

Any caller that compiles a callable which can raise inherits a silent wrong-result on retry — the failure is masked as an empty-output success. A binding can work around it by poisoning the compiled wrapper after a first-trace failure, but the clean fix is to not mark the entry filled until it actually is.



Discussion:
> * set `entry.empty = false` only **after** `compile_trace` (and the dfs / simplify / fuse) complete successfully

This looks a nice and simple fix.



Patch:
diff --git a/mlx/compile.cpp b/mlx/compile.cpp
index 6f3d89cf28..12d7397be4 100644
--- a/mlx/compile.cpp
+++ b/mlx/compile.cpp
@@ -1124,8 +1124,6 @@ ArrayFnWithExtra compile(
 
     // No matching cache entry existed, so compile
     if (entry.empty) {
-      // Mark the entry as not empty since we are about to fill it
-      entry.empty = false;
       // Set the constants
       entry.constants = std::move(constants);
       // Trace to build the graph
@@ -1151,6 +1149,11 @@ ArrayFnWithExtra compile(
       if (mode != CompileMode::no_fuse) {
         compile_fuse(entry.tape, parents_map, entry.inputs, entry.outputs);
       }
+
+      // Mark the entry as filled only after every step above completed, so
+      // a throwing first trace leaves the entry empty and a later call
+      // re-traces cleanly instead of hitting a half-filled entry
+      entry.empty = false;
     }
 
     // At this point we must have a tape, now replace the placeholders


Test patch:
diff --git a/tests/compile_tests.cpp b/tests/compile_tests.cpp
index d2146f3bad..2311e99098 100644
--- a/tests/compile_tests.cpp
+++ b/tests/compile_tests.cpp
@@ -816,3 +816,31 @@ TEST_CASE("test compile random bits") {
   auto out = compile(fun)({in})[0];
   CHECK(array_equal(out, expected).item<bool>());
 }
+
+TEST_CASE("test compile throwing first trace does not poison cache") {
+  // A nullary function whose first trace raises. Without rolling the cache
+  // entry back, the second call with matching (empty) inputs would hit a
+  // half-filled entry, skip tracing, and silently return empty outputs.
+  bool should_throw = true;
+  auto fun = [&should_throw](const std::vector<array>& inputs) {
+    if (should_throw) {
+      throw std::runtime_error("trace failure");
+    }
+    return std::vector<array>{array(1.0f) + array(2.0f)};
+  };
+
+  auto cfun = compile(fun);
+
+  // First call: the trace raises and must propagate.
+  CHECK_THROWS_AS(cfun({}), std::runtime_error);
+
+  // Second call: still raising. It must re-trace and raise again rather than
+  // return an empty result from a poisoned cache entry.
+  CHECK_THROWS_AS(cfun({}), std::runtime_error);
+
+  // Once the function stops raising, a retry must trace cleanly and succeed.
+  should_throw = false;
+  auto out = cfun({});
+  REQUIRE_EQ(out.size(), 1);
+  CHECK_EQ(out[0].item<float>(), 3.0f);
+}