fix: run compiled functions on the task-scoped default device

Source/MLX/Transforms+Compile.swift

@@ -94,6 +94,19 @@ final class CompiledFunction: @unchecked (Sendable) {
             evalLock.unlock()
         }
 
+        // The backend keys its compile cache on, and builds the compiled
+        // primitive for, the C++ default stream, which never sees the
+        // task-scoped device from `Device.withDefaultDevice`. Without this, a
+        // function first traced on the GPU is reused inside a CPU-scoped call
+        // and runs its fused kernel on the GPU against CPU-stream inputs.
+        var previousDevice = mlx_device_new()
+        mlx_get_default_device(&previousDevice)
+        mlx_set_default_device(Device.defaultDevice().ctx)
+        defer {
+            mlx_set_default_device(previousDevice)
+            mlx_device_free(previousDevice)
+        }
+
         let innerInputs = arguments + stateInputs
         let innerInputsVector = new_mlx_vector_array(innerInputs)
         defer { mlx_vector_array_free(innerInputsVector) }

Tests/MLXTests/TransformTests.swift

@@ -133,6 +133,36 @@ class TransformTests: XCTestCase {
         assertEqual(r1, r3)
     }
 
+    func testCompileHonorsScopedDefaultDevice() {
+        // A function traced on the GPU must not be reused inside a CPU-scoped
+        // call: the backend caches compiled graphs per default stream, so the
+        // CPU call has to retrace for its own stream.
+        var traces = 0
+        let compiled = compile { (inputs: [MLXArray]) -> [MLXArray] in
+            traces += 1
+            return [square(inputs[0] * inputs[1])]
+        }
+
+        let i1 = MLXRandom.normal([20, 20])
+        let i2 = MLXRandom.normal([20, 20])
+        let expected = square(i1 * i2)
+
+        assertEqual(compiled([i1, i2])[0], expected)
+        XCTAssertEqual(traces, 1)
+
+        let cpuResult = Device.withDefaultDevice(.cpu) {
+            let r = compiled([i1, i2])[0]
+            eval(r)
+            return r
+        }
+        assertEqual(cpuResult, expected)
+        XCTAssertEqual(traces, 2)
+
+        // Back on the GPU, the original entry is reused.
+        assertEqual(compiled([i1, i2])[0], expected)
+        XCTAssertEqual(traces, 2)
+    }
+
     class CompileTestState: CustomStringConvertible, Updatable {
         var y: MLXArray
         var o: MLXArray?