From 82ee95fd1f83a65c6e2ab9f5d59279de8affbdea Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Thu, 8 Oct 2026 19:05:02 -0700
Subject: [PATCH 1/4] be able to trigger workflows on main.  notes about
 upcoming releases (#489)

* be able to trigger workflows on main.  notes about upcoming releases
Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com>
---
 .github/workflows/pull_request.yml | 6 +++++-
 README.md                          | 4 ++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
index 7dcd7ae49..f98c766b8 100644
--- a/.github/workflows/pull_request.yml
+++ b/.github/workflows/pull_request.yml
@@ -1,6 +1,8 @@
 name: Build and Test
 
-on: pull_request
+on:
+  pull_request:
+  workflow_dispatch:
 
 permissions:
   contents: read
@@ -157,6 +159,7 @@ jobs:
   linux_build_cmake_container:
     name: Linux Container CMake Swift Build (${{ matrix.container }} ${{ matrix.arch }})
     needs: lint
+    if: github.repository == 'ml-explore/mlx-swift'
     strategy:
       fail-fast: false
       matrix:
@@ -194,6 +197,7 @@ jobs:
   linux_build_cmake_cuda:
     name: Linux CUDA CMake Swift Build (cuda-12.9, swift-6.2.3-RELEASE, ubuntu-24.04, x86_64)
     needs: lint
+    if: github.repository == 'ml-explore/mlx-swift'
     runs-on: "gpu-t4-4-core"
     env:
       # Note 12/18/2025: the CI machine does not meet CUDA 13's driver requirement.
diff --git a/README.md b/README.md
index 540607ecd..035e09b15 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,8 @@ MLX is an array framework for machine learning on Apple silicon. MLX Swift
 expands MLX to the Swift language, making research and experimentation easier
 on Apple silicon.
 
+[View Current Release Planning](https://github.com/ml-explore/mlx-swift/issues?q=is%3Aissue+is%3Aopen+label%3Arelease-planning) and read about [Contributing](#contributing).
+
 ## Language Models
 
 LLM and VLM implementations are available in [mlx-swift-lm](https://github.com/ml-explore/mlx-swift-lm).
@@ -185,6 +187,8 @@ contributors](ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
 to MLX Swift and wish to be acknowledged, please add your name to the list in your
 pull request.
 
+Smaller, focused PRs are easier and quicker to review.  Sometimes larger PRs can't be helped.  If you want to discuss, feel free to open an issue or ask questions in a PR.
+
 MLX Swift was initially developed by David Koski and Ronan Collobert, and is
 now maintained by David Koski. MLX Swift is built on top of
 [MLX](https://github.com/ml-explore/mlx), which was initially developed with

From 4dc3edc4e900f56eea00da55999442d09f6bedf1 Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Thu, 8 Oct 2026 19:05:17 -0700
Subject: [PATCH 2/4] fix xcframework builds (#495)

- give me an easy way to trigger dry runs
- the build failed because JIT wasn't available on TVOS
- this isn't new, so I suspect it was a change on the mlx side that we are now seeing
- the guard wasn't wide enough -- just visionOS and iOS, widen to IPHONE (which covers all iOS based OS targets)
---
 .github/workflows/release.yml                       | 13 ++++++++++++-
 .../Cmlx/mlx-conditional/compiled_conditional.cpp   |  2 +-
 .../mlx-conditional/jit_compiler_conditional.cpp    |  2 +-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 7f8e7e61a..044a21a3e 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -27,6 +27,7 @@ jobs:
         run: xcodebuild -showComponent MetalToolchain
 
       - name: Setup gh
+        if: github.event.release.tag_name || inputs.tag
         shell: sh
         run: |
           brew install gh
@@ -44,10 +45,19 @@ jobs:
         shell: bash
         working-directory: xcode/build/output
         run: |
-          for framework in MLX Cmlx MLXNN MLXOptimizers; do
+          for framework in MLX Cmlx cLogSupport MLXNN MLXOptimizers; do
             zip -r -y "${framework}.xcframework.zip" "${framework}.xcframework"
           done
 
+      - name: Upload XCFrameworks as workflow artifacts
+        if: ${{ !(github.event.release.tag_name || inputs.tag) }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: xcframeworks
+          if-no-files-found: error
+          retention-days: 7
+          path: xcode/build/output/*.xcframework.zip
+
       - name: Upload XCFrameworks to release
         if: github.event.release.tag_name || inputs.tag
         env:
@@ -56,5 +66,6 @@ jobs:
           gh release upload "${{ github.event.release.tag_name || inputs.tag }}" \
             xcode/build/output/MLX.xcframework.zip \
             xcode/build/output/Cmlx.xcframework.zip \
+            xcode/build/output/cLogSupport.xcframework.zip \
             xcode/build/output/MLXNN.xcframework.zip \
             xcode/build/output/MLXOptimizers.xcframework.zip
diff --git a/Source/Cmlx/mlx-conditional/compiled_conditional.cpp b/Source/Cmlx/mlx-conditional/compiled_conditional.cpp
index 8f970cb23..fb0e29c95 100644
--- a/Source/Cmlx/mlx-conditional/compiled_conditional.cpp
+++ b/Source/Cmlx/mlx-conditional/compiled_conditional.cpp
@@ -3,7 +3,7 @@
 #include <TargetConditionals.h>
 
 // select the correct cpu compile system based on TARGET_OS
-#if TARGET_OS_IOS || TARGET_OS_VISION
+#if TARGET_OS_IPHONE
 #include "../mlx/mlx/backend/no_cpu/compiled.cpp"
 #else
 #include "../mlx/mlx/backend/cpu/compiled.cpp"
diff --git a/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp b/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp
index 91e81980f..47602719a 100644
--- a/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp
+++ b/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp
@@ -7,6 +7,6 @@
 // `compiled_conditional.cpp` selects `no_cpu/compiled.cpp` on those platforms,
 // so it isn't needed there. Keep this condition in sync with
 // `compiled_conditional.cpp`.
-#if !(TARGET_OS_IOS || TARGET_OS_VISION)
+#if !(TARGET_OS_IPHONE)
 #include "../mlx/mlx/backend/cpu/jit_compiler.cpp"
 #endif

From 1e61d9c29255443b54c712bb8c226554127e960b Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:51:50 -0700
Subject: [PATCH 3/4] fix lock inversion between evalLock and dispatch_once on
 globalStreams (#499)

- globalStreams could take the evalLock
- evalLock could take the globalStreams lock (dispatch_once)
- break the cycle by making the no-args init lock free
---
 Source/MLX/Stream.swift                       | 13 ++--
 .../MLXTests/StreamInitializationTests.swift  | 65 +++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 Tests/MLXTests/StreamInitializationTests.swift

diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift
index 8b31b4105..e8f5aa761 100644
--- a/Source/MLX/Stream.swift
+++ b/Source/MLX/Stream.swift
@@ -166,8 +166,9 @@ private final class RealizedStreamPromise: StreamPromise, Sendable {
 
 /// TaskLocal promise for a pair of Streams -- one for CPU and one for GPU.
 ///
-/// Implementation note: the default device will be a realized Stream and the
-/// none default device will be a lazy promise.
+/// Implementation note: in the scoped promises the default device will be a
+/// realized Stream and the non-default device will be a lazy promise.  The
+/// global promise is lazy for both, see ``init()``.
 private final class StreamPairPromise: Sendable {
     let cpu: StreamPromise
     let gpu: StreamPromise
@@ -181,6 +182,10 @@ private final class StreamPairPromise: Sendable {
     }
 
     /// Initialize the global promise.
+    ///
+    /// Both streams are lazy here to avoid the lock inversion
+    /// from `let globalStreams = StreamPairPromise()` and evalLock.
+    /// See StreamInitializationTests.
     init() {
         var default_ctx = mlx_device_new()
         mlx_get_default_device(&default_ctx)
@@ -188,11 +193,11 @@ private final class StreamPairPromise: Sendable {
         let defaultDevice = Device(default_ctx)
         self.defaultDeviceType = defaultDevice.deviceType
         if defaultDevice.deviceType == .cpu {
-            self.cpu = RealizedStreamPromise(defaultDevice)
+            self.cpu = LazyStreamPromise(defaultDevice)
             self.gpu = LazyStreamPromise(Device(.gpu))
         } else {
             self.cpu = LazyStreamPromise(Device(.cpu))
-            self.gpu = RealizedStreamPromise(defaultDevice)
+            self.gpu = LazyStreamPromise(defaultDevice)
         }
     }
 
diff --git a/Tests/MLXTests/StreamInitializationTests.swift b/Tests/MLXTests/StreamInitializationTests.swift
new file mode 100644
index 000000000..7ca7b4b30
--- /dev/null
+++ b/Tests/MLXTests/StreamInitializationTests.swift
@@ -0,0 +1,65 @@
+// Copyright © 2026 Apple Inc.
+
+import Foundation
+import Testing
+
+@testable import MLX
+
+/// Body of the exit test: synchronous (it blocks on semaphores) and capture free.
+private func raceFirstUseOfDefaultStreamAgainstEvalLock() {
+    let holdingEvalLock = DispatchSemaphore(value: 0)
+    let finished = DispatchGroup()
+
+    // B: hold evalLock, give A time to enter the once and block on
+    // evalLock, then touch the default stream ourselves
+    finished.enter()
+    Thread {
+        withEvalLock {
+            holdingEvalLock.signal()
+            Thread.sleep(forTimeInterval: 0.5)
+            _ = StreamOrDevice.default.stream
+        }
+        finished.leave()
+    }.start()
+
+    // A: first touch of the default stream while B holds evalLock
+    finished.enter()
+    Thread {
+        holdingEvalLock.wait()
+        _ = Stream.defaultStream
+        finished.leave()
+    }.start()
+
+    if finished.wait(timeout: .now() + 10) == .timedOut {
+        // deadlocked: bail out without running exit handlers, which
+        // could themselves block on the stuck threads
+        FileHandle.standardError.write(
+            Data("deadlock: first use of the default stream vs evalLock\n".utf8))
+        _exit(1)
+    }
+}
+
+/// First use of the global streams must not deadlock against `evalLock`.
+///
+/// The global stream pair is a lazily initialized `static let`, so its
+/// initializer runs under a `dispatch_once`.  If that initializer takes
+/// `evalLock` (to create the `mlx_stream`), there is a lock-order inversion:
+///
+/// - thread A: first touch of the default stream -> inside the once ->
+///   waits for `evalLock`
+/// - thread B: holds `evalLock` (an eval, or a `grad` / `compile` trace) ->
+///   first touch of the default stream (e.g. `StreamOrDevice.default` while
+///   building an op) -> waits for the once
+///
+/// This can only happen on the very first touch in a process, so it runs as
+/// an exit test: the body runs in a fresh child process where the globals
+/// have not been initialized yet.
+@Suite("Stream initialization")
+struct StreamInitializationTests {
+
+    @Test func firstUseOfDefaultStreamWhileAnotherThreadHoldsEvalLock() async {
+        await #expect(processExitsWith: .success) {
+            raceFirstUseOfDefaultStreamAgainstEvalLock()
+        }
+    }
+}

From 2f96983a003785234974cc2b0b73138bef7dcf43 Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Thu, 15 Oct 2026 09:36:00 -0700
Subject: [PATCH 4/4] idea for doing synchronous evals but _not_ holding
 evalLock (#474)

* idea for doing synchronous evals but _not_ holding evalLock

- do the synchronous mutation under lock with asyncEval
- synchronous eval then does not need lock

* PR feedback: some evaluation under transform is allowed, e.g. in grad()
---
 Source/MLX/Export.swift                    |  14 +-
 Source/MLX/MLXArray.swift                  |   4 +-
 Source/MLX/Transforms+Eval.swift           | 170 +++--
 Tests/MLXTests/EvalInTransformsTests.swift | 369 ++++++++++
 Tests/MLXTests/EvalTests.swift             | 767 +++++++++++++++++++++
 skills/mlx-swift/references/concurrency.md |  55 +-
 6 files changed, 1308 insertions(+), 71 deletions(-)
 create mode 100644 Tests/MLXTests/EvalInTransformsTests.swift
 create mode 100644 Tests/MLXTests/EvalTests.swift

diff --git a/Source/MLX/Export.swift b/Source/MLX/Export.swift
index e12274c63..9219d2530 100644
--- a/Source/MLX/Export.swift
+++ b/Source/MLX/Export.swift
@@ -144,8 +144,11 @@ public final class FunctionExporterSingle {
         let closure = new_mlx_kwargs_closure(keys: keys, f)
         defer { mlx_closure_kwargs_free(closure) }
 
+        // exporting traces `f`, so it needs evalLock like the other transforms
         _ = try withError {
-            mlx_export_function_kwargs(url.path, closure, positionalArgs, kwargs, shapeless)
+            withEvalLock {
+                mlx_export_function_kwargs(url.path, closure, positionalArgs, kwargs, shapeless)
+            }
         }
     }
 }
@@ -204,8 +207,11 @@ public final class FunctionExporterMultiple {
                 uniquingKeysWith: { a, b in a }))
         defer { mlx_map_string_to_array_free(kwargs) }
 
+        // each call traces `f`, so it needs evalLock like the other transforms
         _ = try withError {
-            mlx_function_exporter_apply_kwargs(exporter, positionalArgs, kwargs)
+            withEvalLock {
+                mlx_function_exporter_apply_kwargs(exporter, positionalArgs, kwargs)
+            }
         }
     }
 }
@@ -292,7 +298,9 @@ public final class ImportedFunction {
         defer { mlx_map_string_to_array_free(kwargs) }
 
         _ = try withError {
-            mlx_imported_function_apply_kwargs(&result, ctx, positionalArgs, kwargs)
+            withEvalLock {
+                mlx_imported_function_apply_kwargs(&result, ctx, positionalArgs, kwargs)
+            }
         }
 
         return mlx_vector_array_values(result)
diff --git a/Source/MLX/MLXArray.swift b/Source/MLX/MLXArray.swift
index 46b260e35..8570cbacd 100644
--- a/Source/MLX/MLXArray.swift
+++ b/Source/MLX/MLXArray.swift
@@ -554,9 +554,7 @@ public final class MLXArray {
     /// MLX is lazy and arrays are not fully realized until they are evaluated.  This method is typically
     /// not needed as all reads ensure the contents are evaluated.
     public func eval() {
-        _ = withEvalLock {
-            mlx_array_eval(ctx)
-        }
+        MLX.eval(self)
     }
 
     /// Replace the contents with a reference to a new array (INTERNAL).
diff --git a/Source/MLX/Transforms+Eval.swift b/Source/MLX/Transforms+Eval.swift
index 64e76098a..ac7b430f0 100644
--- a/Source/MLX/Transforms+Eval.swift
+++ b/Source/MLX/Transforms+Eval.swift
@@ -7,46 +7,55 @@ import Foundation
 /// a recursive lock to handle any cases where a closure might
 /// call back into eval.
 ///
-/// Acquire it through ``withEvalLock(_:)`` rather than directly: debug builds
-/// record ownership there so that the `evalLock` → `CompiledFunction.lock`
-/// ordering can be checked.
+/// Acquire it through ``withEvalLock(_:)`` or ``withEvalLockState(_:)`` rather
+/// than directly: ownership is recorded there, both so that
+/// ``eval(_:)-(Collection<MLXArray>)`` can tell whether it may release the lock
+/// while waiting, and (in debug builds) so that the `evalLock` →
+/// `CompiledFunction.lock` ordering can be checked.
 ///
 /// `evalLock` is the outermost lock in the library.  Nothing may be acquired
 /// before it and then be waited on from inside it; see `CompiledFunction.call`.
 let evalLock = NSRecursiveLock()
 
-#if DEBUG
+/// Per-thread `evalLock` ownership.
+///
+/// Explicit bookkeeping is required: `evalLock` is recursive, so a
+/// successful `try()` cannot distinguish "this thread owns it" from "nobody
+/// owns it".  The depth is thread-local, so it needs no lock of its own --
+/// which matters, because this runs while holding `evalLock`.
+///
+/// This is deliberately per *thread* rather than a `@TaskLocal`: a lock is
+/// owned by a thread, and a task-local binding would be inherited by e.g. a
+/// `Task {}` created inside a traced function, which runs on another thread
+/// that does not hold the lock.
+enum EvalLockOwnership {
+
+    private static let depthKey: pthread_key_t = {
+        var key = pthread_key_t()
+        pthread_key_create(&key, nil)
+        return key
+    }()
+
+    static var isHeldByCurrentThread: Bool {
+        depth > 0
+    }
 
-    /// Per-thread `evalLock` ownership, recorded so the ordering invariant in
-    /// `CompiledFunction.call` can be checked in debug builds.
-    ///
-    /// Explicit bookkeeping is required: `evalLock` is recursive, so a
-    /// successful `try()` cannot distinguish "this thread owns it" from "nobody
-    /// owns it".  The depth is thread-local, so it needs no lock of its own --
-    /// which matters, because this runs while holding `evalLock`.
-    enum EvalLockOwnership {
-
-        private static let depthKey: pthread_key_t = {
-            var key = pthread_key_t()
-            pthread_key_create(&key, nil)
-            return key
-        }()
-
-        static var isHeldByCurrentThread: Bool {
-            depth > 0
-        }
+    static var depth: Int {
+        Int(bitPattern: pthread_getspecific(depthKey))
+    }
 
-        static var depth: Int {
-            Int(bitPattern: pthread_getspecific(depthKey))
-        }
+    static func entered() {
+        pthread_setspecific(depthKey, UnsafeRawPointer(bitPattern: depth + 1))
+    }
 
-        static func entered() {
-            pthread_setspecific(depthKey, UnsafeRawPointer(bitPattern: depth + 1))
-        }
+    static func exited() {
+        pthread_setspecific(depthKey, UnsafeRawPointer(bitPattern: depth - 1))
+    }
+}
 
-        static func exited() {
-            pthread_setspecific(depthKey, UnsafeRawPointer(bitPattern: depth - 1))
-        }
+#if DEBUG
+
+    extension EvalLockOwnership {
 
         // Counters, so a test can assert both that the invariant held and that
         // it was actually exercised -- a check count of zero would mean the test
@@ -95,24 +104,57 @@ let evalLock = NSRecursiveLock()
 
 #endif
 
+/// How the current call to ``withEvalLockState(_:)`` holds
+/// ``evalLock``.
+enum EvalLockState {
+    /// This call acquired ``evalLock``: it is released when the body returns.
+    case outermost
+
+    /// The current thread already held ``evalLock`` -- e.g. inside a `grad`,
+    /// `vjp`, `jvp`, `vmap` or `compile` trace, or a nested call.  The lock
+    /// stays held after the body returns.
+    case nested
+}
+
 /// Acquire ``evalLock`` for the duration of `body`.
 ///
-/// This is the single choke point for ``evalLock``; debug builds record
-/// ownership here so that `CompiledFunction` can check that its per-instance
-/// lock is never taken outside it.
+/// This is the single choke point for ``evalLock``; ownership is recorded here
+/// so that callers can tell whether the current thread already held it (see
+/// ``EvalLockState``) and so that `CompiledFunction` can check that its
+/// per-instance lock is never taken outside it.
 @inline(__always)
-func withEvalLock<R>(_ body: () throws -> R) rethrows -> R {
+func withEvalLockState<R>(_ body: (EvalLockState) throws -> R) rethrows -> R {
     evalLock.lock()
-    #if DEBUG
-        EvalLockOwnership.entered()
-    #endif
+    let state: EvalLockState = EvalLockOwnership.isHeldByCurrentThread ? .nested : .outermost
+    EvalLockOwnership.entered()
     defer {
-        #if DEBUG
-            EvalLockOwnership.exited()
-        #endif
+        EvalLockOwnership.exited()
         evalLock.unlock()
     }
-    return try body()
+    return try body(state)
+}
+
+/// Acquire ``evalLock`` for the duration of `body`.
+///
+/// See ``withEvalLockState(_:)``.
+@inline(__always)
+func withEvalLock<R>(_ body: () throws -> R) rethrows -> R {
+    try withEvalLockState { _ in try body() }
+}
+
+/// Is the array's data computed?
+///
+/// This is the observable difference between ``eval(_:)-(MLXArray...)`` and
+/// ``asyncEval(_:)-(Collection<MLXArray>)``, and the only direct way to check
+/// that `eval` is still synchronous now that it schedules with
+/// `mlx_async_eval` and waits with `mlx_eval`.
+///
+/// Internal: exists for the tests (`EvalTests`), which cannot reach `Cmlx`
+/// directly.  Call this only on the thread that owns `array`.
+func isEvaluated(_ array: MLXArray) -> Bool {
+    var available = false
+    _mlx_array_is_available(&available, array.ctx)
+    return available
 }
 
 /// Evaluate one or more `MLXArray`
@@ -120,23 +162,49 @@ func withEvalLock<R>(_ body: () throws -> R) rethrows -> R {
 /// ### See Also
 /// - <doc:lazy-evaluation>
 public func eval(_ arrays: MLXArray...) {
-    let vector_array = new_mlx_vector_array(arrays)
-    _ = withEvalLock {
-        mlx_eval(vector_array)
-    }
-    mlx_vector_array_free(vector_array)
+    eval(arrays)
 }
 
 /// Evaluate one or more `MLXArray`
 ///
+/// This is synchronous: the arrays are computed when it returns.  `evalLock`
+/// is only held while the work is scheduled, not while waiting for it to
+/// finish, so evaluations on other threads can overlap.
+///
+/// When the current thread already holds `evalLock` -- inside a `grad`,
+/// `vjp`, `jvp`, `vmap` or `compile` trace -- the lock cannot be released for
+/// the wait, and evaluation is fully synchronous under the lock.  That path is
+/// also required for correctness: evaluating an intermediate inside `grad` or
+/// `vjp` is allowed, but scheduling it asynchronously is not.
+///
+/// As with any other use of `MLXArray`, evaluation is not thread safe for a
+/// given array: arrays shared between threads (e.g. model weights) should be
+/// evaluated once, on one thread, before they are shared.
+///
 /// ### See Also
 /// - <doc:lazy-evaluation>
 public func eval(_ arrays: some Collection<MLXArray>) {
-    let vector_array = new_mlx_vector_array(arrays)
-    _ = withEvalLock {
-        mlx_eval(vector_array)
+    let vectorArray = new_mlx_vector_array(arrays)
+    defer { mlx_vector_array_free(vectorArray) }
+
+    let waitOutsideLock = withEvalLockState { state in
+        switch state {
+        case .nested:
+            // cannot release the lock to wait, and async eval rejects
+            // vjp/jvp tracers
+            _ = mlx_eval(vectorArray)
+            return false
+
+        case .outermost:
+            // schedule under the lock, wait (below) outside it
+            return mlx_async_eval(vectorArray) == 0
+        }
+    }
+
+    if waitOutsideLock {
+        // the arrays are scheduled, so this only waits for them
+        _ = mlx_eval(vectorArray)
     }
-    mlx_vector_array_free(vector_array)
 }
 
 /// Evaluate one or more `MLXArray` asynchronously.
diff --git a/Tests/MLXTests/EvalInTransformsTests.swift b/Tests/MLXTests/EvalInTransformsTests.swift
new file mode 100644
index 000000000..10473334d
--- /dev/null
+++ b/Tests/MLXTests/EvalInTransformsTests.swift
@@ -0,0 +1,369 @@
+// Copyright © 2026 Apple Inc.
+
+import Foundation
+import MLXNN
+import Testing
+
+@testable import MLX
+
+// MARK: - Support
+
+/// Results of evaluating inside a function transformation.
+///
+/// MLX core (`eval_impl` in `mlx/transforms.cpp`) distinguishes:
+///
+/// - `vjp` / `jvp` (and so `grad`, `valueAndGrad`, custom VJPs): tracers are
+///   real arrays with primitives, so a *synchronous* eval of an intermediate is
+///   allowed and the graph is retained for the backward pass.  An
+///   *asynchronous* eval of a tracer is rejected.  See `test_eval_in_grad` and
+///   `test_async_eval_in_trace` in the python tests.
+/// - `compile` / `vmap`: tracing runs on placeholder inputs with no primitive,
+///   so evaluating anything derived from them is rejected by either kind of
+///   eval.
+///
+/// Several of the failure modes here abort the process (the default error
+/// handler calls `fatalError`) or crash outright (reading the data pointer of
+/// an array whose eval failed), so those cases run as exit tests: the body
+/// runs in a child process and the parent checks how it exited.
+
+/// d/dx sum(x^2) at x = 3, used by most of the `grad` cases.
+private let x0: Float = 3
+private let expectedGradient: Float = 6
+
+private final class Model: Module, UnaryLayer {
+    let linear: Linear
+
+    override init() {
+        linear = Linear(weight: MLXArray([Float(2)], [1, 1]), bias: MLXArray([Float(1)]))
+    }
+
+    func callAsFunction(_ x: MLXArray) -> MLXArray { linear(x) }
+}
+
+// MARK: Bodies (also run inside exit tests, so they must not capture state)
+
+/// `eval` of an intermediate inside `grad`, default error handler.
+private func evalIntermediateInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let y = x * x
+        eval(y)
+        return y.sum()
+    }
+    let d = g(MLXArray(x0))
+    #expect(d.item(Float.self) == expectedGradient)
+}
+
+/// `MLXArray.eval()` of an intermediate inside `grad`.
+private func arrayEvalInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let y = x * x
+        y.eval()
+        return y.sum()
+    }
+    let d = g(MLXArray(x0))
+    #expect(d.item(Float.self) == expectedGradient)
+}
+
+/// `asArray` (debug printing an intermediate) inside `grad`.
+private func asArrayInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let y = x * x
+        #expect(y.asArray(Float.self) == [x0 * x0])
+        return y.sum()
+    }
+    let d = g(MLXArray(x0))
+    #expect(d.item(Float.self) == expectedGradient)
+}
+
+/// `asArray` inside `grad`, under `withError` -- the error (if any) is
+/// captured, so the eval is skipped and `asArray` reads a null pointer.
+private func asArrayInsideGradWithError() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let y = x * x
+        let values = try? withError { y.asArray(Float.self) }
+        #expect(values == [x0 * x0])
+        return y.sum()
+    }
+    let d = g(MLXArray(x0))
+    #expect(d.item(Float.self) == expectedGradient)
+}
+
+/// Control flow on an evaluated value inside `grad` -- the `y[0] >= 0` shape
+/// of python code, which needs a real read of the data.
+private func controlFlowInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let y = x * x
+        if (y .> 0).item(Bool.self) {
+            return y.sum()
+        } else {
+            return (-y).sum()
+        }
+    }
+    let d = g(MLXArray(x0))
+    #expect(d.item(Float.self) == expectedGradient)
+}
+
+/// `eval` inside `valueAndGrad(model:)`'s loss function, e.g. to log the loss.
+private func evalInsideModelValueAndGrad() {
+    let model = Model()
+    eval(model)
+
+    func loss(model: Model, x: MLXArray, y: MLXArray) -> MLXArray {
+        let l = mseLoss(predictions: model(x), targets: y, reduction: .mean)
+        eval(l)
+        _ = l.asArray(Float.self)
+        return l
+    }
+
+    let lg = valueAndGrad(model: model, loss)
+
+    // prediction = 2 * 1 + 1 = 3, target 0: loss = 9,
+    // dL/dw = 2 * 3 * x = 6, dL/db = 2 * 3 = 6
+    let (l, grads) = lg(model, MLXArray([Float(1)], [1, 1]), MLXArray([Float(0)], [1, 1]))
+    eval(l, grads)
+
+    let flat = Dictionary(uniqueKeysWithValues: grads.flattened())
+    #expect(l.item(Float.self) == 9)
+    #expect(flat["linear.weight"]?.asArray(Float.self) == [6])
+    #expect(flat["linear.bias"]?.asArray(Float.self) == [6])
+}
+
+/// `eval` inside the VJP of a custom function: this runs during the backward
+/// pass, where every array is a tracer because core holds `RetainGraph`.
+private func evalInsideCustomVJP() {
+    let f = CustomFunction {
+        Forward { inputs in [inputs[0] * inputs[0]] }
+        VJP { primals, cotangents in
+            let v = 2 * primals[0] * cotangents[0]
+            eval(v)
+            return [v]
+        }
+    }
+
+    let g = grad { (x: MLXArray) -> MLXArray in f([x])[0].sum() }
+    let d = g(MLXArray(x0))
+    #expect(d.item(Float.self) == expectedGradient)
+}
+
+/// Port of `test_eval_in_grad` from the python tests.
+private func evalInsideVJP() {
+    let arr = MLXArray([Float(1)])
+    let cotangent = MLXArray([Float(1), 1])
+    let y = MLXArray([Float(2), 2])
+
+    // reading a value derived from the primal
+    let (_, vjps1) = vjp(
+        { x in
+            let x = x[0] + y
+            let condition = x .< 1
+            _ = condition.asArray(Bool.self)
+            return [x ** 2]
+        }, primals: [arr], cotangents: [cotangent])
+    #expect(vjps1[0].item(Float.self) == 12)
+
+    // evaluating a value derived from the primal
+    let (_, vjps2) = vjp(
+        { x in
+            let x = x[0] + MLXArray([Float(1), 1])
+            eval(x)
+            return [x ** 2]
+        }, primals: [arr], cotangents: [cotangent])
+    #expect(vjps2[0].item(Float.self) == 8)
+}
+
+/// `eval` of an intermediate inside `jvp`.
+private func evalInsideJVP() {
+    let (_, jvps) = jvp(
+        { x in
+            let y = x[0] * x[0]
+            eval(y)
+            return [y]
+        }, primals: [MLXArray(x0)], tangents: [MLXArray(Float(1))])
+    #expect(jvps[0].item(Float.self) == expectedGradient)
+}
+
+/// `asyncEval` of an intermediate inside `grad` -- not allowed.
+private func asyncEvalInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let y = x * x
+        asyncEval(y)
+        return y.sum()
+    }
+    _ = g(MLXArray(x0))
+}
+
+/// `eval` of a traced value inside `vmap` -- not allowed.
+private func evalInsideVmap() {
+    let mapped = vmap { (x: MLXArray) -> MLXArray in
+        let y = x * 2
+        eval(y)
+        return y
+    }
+    _ = mapped(MLXArray(0 ..< 4).asType(.float32))
+}
+
+/// `eval` of a traced value inside `compile` -- not allowed.
+private func evalInsideCompile() {
+    let compiled = compile { (x: MLXArray) -> MLXArray in
+        let y = x * 2
+        eval(y)
+        return y
+    }
+    _ = compiled(MLXArray(Float(3)))
+}
+
+// MARK: - Tests
+
+@Suite("eval inside function transformations")
+struct EvalInTransformsTests {
+
+    // MARK: Allowed: synchronous eval of vjp/jvp tracers
+
+    // These run as exit tests so that a regression (fatalError from the default
+    // error handler, or a crash reading an unevaluated array) fails the test
+    // rather than taking down the test runner.
+
+    @Test func evalInsideGrad() async {
+        await #expect(processExitsWith: .success) {
+            evalIntermediateInsideGrad()
+        }
+    }
+
+    @Test func arrayEvalMethodInsideGrad() async {
+        await #expect(processExitsWith: .success) {
+            arrayEvalInsideGrad()
+        }
+    }
+
+    @Test func asArrayInsideGrad() async {
+        await #expect(processExitsWith: .success) {
+            MLXTests.asArrayInsideGrad()
+        }
+    }
+
+    @Test func asArrayInsideGradWithError() async {
+        await #expect(processExitsWith: .success) {
+            MLXTests.asArrayInsideGradWithError()
+        }
+    }
+
+    @Test func controlFlowOnAValueInsideGrad() async {
+        await #expect(processExitsWith: .success) {
+            controlFlowInsideGrad()
+        }
+    }
+
+    @Test func evalInsideModelValueAndGrad() async {
+        await #expect(processExitsWith: .success) {
+            MLXTests.evalInsideModelValueAndGrad()
+        }
+    }
+
+    @Test func evalInsideCustomVJP() async {
+        await #expect(processExitsWith: .success) {
+            MLXTests.evalInsideCustomVJP()
+        }
+    }
+
+    @Test func evalInsideVJP() async {
+        await #expect(processExitsWith: .success) {
+            MLXTests.evalInsideVJP()
+        }
+    }
+
+    @Test func evalInsideJVP() async {
+        await #expect(processExitsWith: .success) {
+            MLXTests.evalInsideJVP()
+        }
+    }
+
+    /// In-process and checked, so a regression reports MLX's error message.
+    @Test func checkedEvalInsideGrad() throws {
+        var thrown: Error?
+        let g = grad { (x: MLXArray) -> MLXArray in
+            let y = x * x
+            do {
+                try checkedEval(y)
+            } catch {
+                thrown = error
+            }
+            return y.sum()
+        }
+        let d = g(MLXArray(x0))
+
+        #expect(thrown == nil, "\(String(describing: thrown))")
+        #expect(d.item(Float.self) == expectedGradient)
+    }
+
+    /// In-process and checked, inside `valueAndGrad(model:)`.
+    @Test func checkedEvalInsideModelValueAndGrad() throws {
+        let model = Model()
+        eval(model)
+
+        var thrown: Error?
+        func loss(model: Model, x: MLXArray, y: MLXArray) -> MLXArray {
+            let l = mseLoss(predictions: model(x), targets: y, reduction: .mean)
+            do {
+                try checkedEval(l)
+            } catch {
+                thrown = error
+            }
+            return l
+        }
+
+        let lg = valueAndGrad(model: model, loss)
+        let (l, _) = lg(model, MLXArray([Float(1)], [1, 1]), MLXArray([Float(0)], [1, 1]))
+
+        #expect(thrown == nil, "\(String(describing: thrown))")
+        #expect(l.item(Float.self) == 9)
+    }
+
+    // MARK: Not allowed
+
+    @Test func asyncEvalInsideGradAborts() async {
+        await #expect(processExitsWith: .failure) {
+            asyncEvalInsideGrad()
+        }
+    }
+
+    @Test func checkedAsyncEvalInsideGradThrows() {
+        #expect {
+            try withError { asyncEvalInsideGrad() }
+        } throws: { error in
+            String(describing: error).contains(
+                "[async_eval] Not allowed inside a graph transformation")
+        }
+    }
+
+    @Test func evalInsideVmapAborts() async {
+        await #expect(processExitsWith: .failure) {
+            evalInsideVmap()
+        }
+    }
+
+    @Test func checkedEvalInsideVmapThrows() {
+        #expect {
+            try withError { evalInsideVmap() }
+        } throws: { error in
+            String(describing: error).contains(
+                "[eval] Attempting to eval an array during function transformations like compile or vmap"
+            )
+        }
+    }
+
+    @Test func evalInsideCompileAborts() async {
+        await #expect(processExitsWith: .failure) {
+            evalInsideCompile()
+        }
+    }
+
+    @Test func checkedEvalInsideCompileThrows() {
+        #expect {
+            try withError { evalInsideCompile() }
+        } throws: { error in
+            String(describing: error).contains(
+                "[eval] Attempting to eval an array during function transformations like compile or vmap"
+            )
+        }
+    }
+}
diff --git a/Tests/MLXTests/EvalTests.swift b/Tests/MLXTests/EvalTests.swift
new file mode 100644
index 000000000..abaf825dd
--- /dev/null
+++ b/Tests/MLXTests/EvalTests.swift
@@ -0,0 +1,767 @@
+// Copyright © 2026 Apple Inc.
+
+import Foundation
+import MLXNN
+import MLXOptimizers
+import Testing
+
+@testable import MLX
+
+// MARK: - Support
+
+/// Run `body` once per thread on `count` real (non-cooperative) threads, all
+/// released from a rendezvous so the calls actually overlap, and wait for them.
+///
+/// Real `Thread`s rather than a `TaskGroup`: `eval` blocks, and the point of
+/// these tests is to have several threads blocked inside the eval machinery at
+/// the same time.  Blocking the cooperative pool would make the amount of
+/// overlap depend on the pool width instead of on the code under test.
+///
+/// - Returns: `true` if every thread finished within `timeout`.
+@discardableResult
+private func onThreads(
+    _ count: Int, timeout: TimeInterval = 300, _ body: @escaping @Sendable (Int) -> Void
+) -> Bool {
+    // `ready` is a rendezvous rather than a completion group: it is entered
+    // once per thread up front, and each thread leaves its own entry and then
+    // waits for the group to drain -- so nobody starts until everybody has
+    // arrived.  `finished` is the ordinary completion use.
+    let ready = DispatchGroup()
+    let finished = DispatchGroup()
+
+    for _ in 0 ..< count {
+        ready.enter()
+        finished.enter()
+    }
+
+    for i in 0 ..< count {
+        let thread = Thread {
+            ready.leave()
+            // bounded, so a thread that never starts cannot wedge the others
+            _ = ready.wait(timeout: .now() + 60)
+            body(i)
+            finished.leave()
+        }
+        thread.name = "EvalTests-\(i)"
+        thread.stackSize = 4 << 20
+        thread.start()
+    }
+
+    return finished.wait(timeout: .now() + timeout) == .success
+}
+
+/// Lock protected collection point for values produced on worker threads.
+///
+/// Only `Sendable` scalars are collected -- never an `MLXArray`.  `MLXArray` is
+/// a reference type wrapping an `mlx_array` handle and is *not* thread safe:
+/// `evalLock` serializes evaluation, not array lifetime or mutation.  Every
+/// test below therefore builds and reads its arrays on a single thread, and
+/// shares only the resulting numbers.
+private final class Collected<Element: Sendable>: @unchecked Sendable {
+    private let lock = NSLock()
+    private var values = [Element]()
+
+    func append(_ value: Element) {
+        lock.withLock { values.append(value) }
+    }
+
+    var all: [Element] {
+        lock.withLock { values }
+    }
+}
+
+/// A `Bool` that can be set on one thread and read on another.
+private final class Flag: @unchecked Sendable {
+    private let lock = NSLock()
+    private var value: Bool
+
+    init(_ value: Bool = false) {
+        self.value = value
+    }
+
+    var isSet: Bool {
+        lock.withLock { value }
+    }
+
+    func set(_ newValue: Bool) {
+        lock.withLock { value = newValue }
+    }
+}
+
+/// Has the array's data been computed?
+///
+/// Thin wrapper over the library's internal `isEvaluated`, reached through
+/// `@testable import`.  Call this only on the thread that owns `array`.
+private func isAvailable(_ array: MLXArray) -> Bool {
+    MLX.isEvaluated(array)
+}
+
+/// A graph that takes a while to evaluate, used to open a window during which
+/// another thread can be observed making progress (or not).
+///
+/// The values are chosen so the result is exact and independent of the number
+/// of matmuls: ones × (ones / n) == ones, so the sum is always `size * size`.
+private func heavyGraph(matmuls: Int, size: Int, stream: StreamOrDevice) -> MLXArray {
+    var x = MLXArray.ones([size, size], stream: stream)
+    let w = MLX.divide(
+        MLXArray.ones([size, size], stream: stream), MLXArray(Float(size)), stream: stream)
+    for _ in 0 ..< matmuls {
+        x = MLX.matmul(x, w, stream: stream)
+    }
+    return x
+}
+
+private func timed(_ body: () -> Void) -> TimeInterval {
+    let start = Date()
+    body()
+    return Date().timeIntervalSince(start)
+}
+
+// MARK: - Tests
+
+/// Behavior of ``eval(_:)-(MLXArray...)``, ``asyncEval(_:)-(Collection<MLXArray>)`` and the
+/// other `evalLock` holders when called from several threads at once.
+///
+/// `eval` schedules under `evalLock` (via `mlx_async_eval`) and then waits
+/// *outside* it (via `mlx_eval`), so two properties need coverage:
+///
+/// - it is still synchronous -- results are available when it returns; and
+/// - it no longer serializes the waiting, so unrelated work on other threads
+///   makes progress while an eval is in flight.
+///
+/// Note what is *not* claimed: an individual `MLXArray` is not thread safe, so
+/// each thread here owns the arrays it creates and only scalar results cross
+/// threads.  What is under test is the shared global state behind eval --
+/// the stream/scheduler, the compiler cache, the tracing state -- which
+/// `evalLock` guards.
+///
+/// The suite is `.serialized` because several tests are stress or timing
+/// shaped: running them against each other would make each one's measurement
+/// depend on the others.
+@Suite("eval + evalLock", .serialized)
+struct EvalTests {
+
+    // MARK: Semantics
+
+    /// `eval` must not return until the arrays are computed, even though the
+    /// wait now happens outside `evalLock`.  `asyncEval` is included as the
+    /// contrast: it is allowed to return early (and normally does, though that
+    /// is a race, so only the "available after eval" direction is asserted).
+    @Test func evalIsSynchronous() {
+        for _ in 0 ..< 20 {
+            let a = heavyGraph(matmuls: 20, size: 256, stream: .default)
+            #expect(!isAvailable(a), "the graph was computed before eval was called")
+            eval(a)
+            #expect(isAvailable(a), "eval returned before the array was available")
+
+            // variadic, collection and Any-collecting overloads all go through
+            // the same schedule-then-wait path
+            let b = heavyGraph(matmuls: 20, size: 256, stream: .default)
+            let c = heavyGraph(matmuls: 20, size: 256, stream: .default)
+            eval(b, c)
+            #expect(isAvailable(b) && isAvailable(c))
+
+            let d = heavyGraph(matmuls: 20, size: 256, stream: .default)
+            eval([d])
+            #expect(isAvailable(d))
+
+            let e = heavyGraph(matmuls: 20, size: 256, stream: .default)
+            eval(["key": e] as [String: MLXArray])
+            #expect(isAvailable(e))
+
+            let f = heavyGraph(matmuls: 20, size: 256, stream: .default)
+            asyncEval(f)
+            eval(f)
+            #expect(isAvailable(f), "eval after asyncEval must still wait")
+        }
+    }
+
+    /// `eval` of an already evaluated array, and repeated evals of the same
+    /// array, must be harmless -- the double schedule (`async_eval` then
+    /// `eval`) must not resubmit or fail on a computed graph.
+    @Test func repeatedEvalIsIdempotent() {
+        let a = MLXArray(0 ..< 8, [2, 4]).asType(.float32) * 2
+        for _ in 0 ..< 10 {
+            eval(a)
+            #expect(isAvailable(a))
+        }
+        #expect(a.sum().item(Float.self) == 56)
+    }
+
+    // MARK: Concurrency, correctness
+
+    /// The baseline: many threads, each building and evaluating its own graph.
+    /// Every thread's result must be its own -- a graph or vector mixed up in
+    /// the shared scheduler state shows up as a wrong value, not just as a
+    /// crash.
+    ///
+    /// Note the shape used by every concurrent test here: worker threads only
+    /// *collect*, and all assertions happen on the test's own thread.  A
+    /// `#expect` evaluated on a raw `Thread` has no current-test context, so a
+    /// failure there would not be attributed to the test.
+    @Test func concurrentEvalOfIndependentGraphs() {
+        let threads = 8
+        let rounds = 50
+        let results = Collected<[Float]>()
+        let available = Collected<Bool>()
+
+        let finished = onThreads(threads) { t in
+            var values = [Float]()
+            for round in 0 ..< rounds {
+                let x = MLXArray(Float(t + 1))
+                let y = MLXArray(Float(round))
+                let r = x * 1000 + y
+                eval(r)
+                available.append(isAvailable(r))
+                values.append(r.item(Float.self))
+            }
+            results.append(values)
+        }
+
+        #expect(finished, "a thread wedged in eval")
+        #expect(available.all.allSatisfy { $0 }, "eval returned before the array was available")
+
+        let all = results.all
+        #expect(all.count == threads)
+        // each thread's values identify it, so a mix-up is visible
+        let expected = Set(
+            (0 ..< threads).map { t in
+                (0 ..< rounds).map { Float((t + 1) * 1000 + $0) }
+            })
+        #expect(Set(all) == expected)
+    }
+
+    /// The same graph shape evaluated on every thread at once.  Unlike the
+    /// test above the threads compute *identical* values from identical
+    /// (separately built) graphs, so all of them contend for the same kernels
+    /// and cache entries while still owning their own arrays.
+    @Test func concurrentEvalOfIdenticalGraphs() {
+        let values = Collected<Float>()
+        let available = Collected<Bool>()
+
+        let finished = onThreads(8) { _ in
+            for _ in 0 ..< 25 {
+                let r = (MLXArray(0 ..< 64, [8, 8]).asType(.float32) * 3).sum()
+                eval(r)
+                available.append(isAvailable(r))
+                values.append(r.item(Float.self))
+            }
+        }
+
+        #expect(finished, "a thread wedged evaluating an identical graph")
+        #expect(values.all.allSatisfy { $0 == 6048 }, "values: \(Set(values.all))")
+        #expect(
+            available.all.allSatisfy { $0 },
+            "eval returned on some thread while its array was still unavailable")
+    }
+
+    /// `asyncEval` (schedule only) and `eval` (schedule then wait) submitted
+    /// from different threads at the same time.  Each thread owns its arrays;
+    /// what is shared is the scheduler they both push work into.
+    @Test func concurrentAsyncEvalAndEval() {
+        let values = Collected<Float>()
+        let available = Collected<Bool>()
+
+        let finished = onThreads(6) { t in
+            for round in 0 ..< 30 {
+                let a = heavyGraph(matmuls: 10, size: 128, stream: .default)
+                let b = a.sum() + Float(round)
+
+                switch t % 3 {
+                case 0:
+                    // schedule only, then wait for it later
+                    asyncEval(b)
+                    eval(b)
+                case 1:
+                    // schedule twice: the second must not lose the first
+                    asyncEval(b)
+                    asyncEval(b)
+                    eval(b)
+                default:
+                    eval(b)
+                }
+
+                available.append(isAvailable(b))
+                values.append(b.item(Float.self) - Float(round))
+            }
+        }
+
+        #expect(finished, "a thread wedged mixing asyncEval and eval")
+        #expect(
+            available.all.allSatisfy { $0 },
+            "eval returned before the array was available")
+        #expect(values.all.count == 6 * 30)
+        #expect(
+            values.all.allSatisfy { $0 == Float(128 * 128) },
+            "values: \(Set(values.all))")
+    }
+
+    /// `MLXArray.eval()`, `item()` and `description` take `evalLock`
+    /// separately from the free `eval` functions; they must mix with it.
+    @Test func concurrentItemDescriptionAndEval() {
+        let items = Collected<Float>()
+        let descriptions = Collected<Int>()
+
+        let finished = onThreads(8) { t in
+            for _ in 0 ..< 40 {
+                let a = MLXArray(0 ..< 16, [4, 4]).asType(.float32)
+                switch t % 4 {
+                case 0:
+                    let r = a.sum()
+                    eval(r)
+                    items.append(r.item(Float.self))
+                case 1:
+                    let r = a.sum()
+                    r.eval()
+                    items.append(r.item(Float.self))
+                case 2:
+                    // item() evaluates implicitly
+                    items.append(a.sum().item(Float.self))
+                default:
+                    // description evaluates under the lock as well
+                    descriptions.append((a * 1).description.count)
+                }
+            }
+        }
+
+        #expect(finished, "a thread wedged in item()/description")
+        #expect(items.all.allSatisfy { $0 == 120 }, "unexpected values: \(Set(items.all))")
+        #expect(descriptions.all.allSatisfy { $0 > 0 })
+    }
+
+    /// Each thread on its own stream.  This also exercises `Stream`'s
+    /// creation, `deinit` and `synchronize`, all of which take `evalLock`, at
+    /// the same time as evals are in flight on other streams.
+    @Test func concurrentEvalAcrossStreams() {
+        let results = Collected<Float>()
+
+        let finished = onThreads(6) { t in
+            for round in 0 ..< 20 {
+                // a fresh stream each round, so Stream.init/deinit run against
+                // the eval traffic as well
+                let stream = Stream(round % 2 == 0 ? Device.gpu : Device.cpu)
+                let s = StreamOrDevice.stream(stream)
+                // MLXArray(_:) is a synchronous copy of host memory -- it has
+                // no stream; only the operation is placed on one
+                let x = MLXArray(Float(t + 1))
+                let r = MLX.multiply(x, MLXArray(Float(7)), stream: s)
+                eval(r)
+                results.append(r.item(Float.self) / Float(t + 1))
+                stream.synchronize()
+            }
+        }
+
+        #expect(finished, "a thread wedged evaluating on its own stream")
+        #expect(results.all.allSatisfy { $0 == 7 }, "values: \(Set(results.all))")
+    }
+
+    /// Compiled functions acquire `evalLock` and then their own per-instance
+    /// lock, so they are the case where lock *ordering* matters.  Hammer them
+    /// against plain evals, including stateful compiled functions whose state
+    /// is evaluated concurrently.
+    @Test func concurrentCompiledCallsAndEval() {
+        #if DEBUG
+            EvalLockOwnership.resetCounters()
+        #endif
+
+        let threads = 6
+        let rounds = 40
+        let results = Collected<[Float]>()
+        let statelessResults = Collected<[Float]>()
+
+        let finished = onThreads(threads) { t in
+            // per-thread state and per-thread compiled function: sharing either
+            // across threads would be a user-level data race on non-Sendable
+            // values, not a lock bug
+            let state = MLXArray(Float(t + 1))
+            func body(_ x: [MLXArray]) -> [MLXArray] {
+                [x[0] * 2 + state]
+            }
+            let compiled = compile(inputs: [state], outputs: [state], body(_:))
+            let stateless = MLX.compile { (x: MLXArray) in x * 3 }
+
+            var values = [Float]()
+            var statelessValues = [Float]()
+            for round in 0 ..< rounds {
+                let r = compiled([MLXArray(Float(round))])
+                eval(r)
+                values.append(r[0].item(Float.self))
+
+                let s = stateless(MLXArray(Float(round)))
+                eval(s, state)
+                statelessValues.append(s.item(Float.self))
+            }
+            results.append(values)
+            statelessResults.append(statelessValues)
+        }
+
+        #expect(finished, "a thread wedged in a compiled call")
+
+        // the stateless function is the same on every thread
+        let statelessExpected = (0 ..< rounds).map { Float($0 * 3) }
+        #expect(
+            statelessResults.all.allSatisfy { $0 == statelessExpected },
+            "a stateless compiled function produced wrong values")
+
+        let all = results.all
+        #expect(all.count == threads)
+        let expected = Set(
+            (0 ..< threads).map { t in
+                (0 ..< rounds).map { Float($0 * 2 + t + 1) }
+            })
+        #expect(Set(all) == expected, "compiled results were mixed up across threads")
+
+        #if DEBUG
+            let counters = EvalLockOwnership.counters
+            #expect(
+                counters.checks > 0,
+                "no lock ordering check was reached, so this test proves nothing")
+            #expect(
+                counters.violations == 0,
+                "an inner lock was taken without evalLock held")
+        #endif
+    }
+
+    /// `grad`/`vjp`/`vmap` trace under `evalLock` while `eval` no longer holds
+    /// it for the whole wait; tracing must still be mutually exclusive, and a
+    /// trace on one thread must not pick up another thread's graph.
+    @Test func concurrentTracingAndEval() {
+        #if DEBUG
+            EvalLockOwnership.resetCounters()
+        #endif
+
+        let gradients = Collected<Float>()
+        let vjps = Collected<Float>()
+        let vmaps = Collected<Float>()
+
+        let finished = onThreads(6) { t in
+            for round in 0 ..< 25 {
+                switch t % 3 {
+                case 0:
+                    // d/dx (x^2) at x = 3 is 6, regardless of what any other
+                    // thread is tracing
+                    let g = MLX.grad { (x: MLXArray) in (x * x).sum() }
+                    let d = g(MLXArray(Float(3)))
+                    eval(d)
+                    gradients.append(d.item(Float.self))
+
+                case 1:
+                    let (_, cotangentProducts) = MLX.vjp(
+                        { (x: [MLXArray]) in [x[0] * 3] },
+                        primals: [MLXArray(Float(round))],
+                        cotangents: [MLXArray(Float(1))])
+                    eval(cotangentProducts)
+                    vjps.append(cotangentProducts[0].item(Float.self))
+
+                default:
+                    let mapped = MLX.vmap { (x: MLXArray) in x * 2 }
+                    let r = mapped(MLXArray(0 ..< 8).asType(.float32))
+                    eval(r)
+                    vmaps.append(r.sum().item(Float.self))
+                }
+            }
+        }
+
+        #expect(finished, "a thread wedged tracing")
+        #expect(gradients.all.allSatisfy { $0 == 6 }, "gradients: \(Set(gradients.all))")
+        #expect(vjps.all.allSatisfy { $0 == 3 }, "vjps: \(Set(vjps.all))")
+        #expect(vmaps.all.allSatisfy { $0 == 56 }, "vmaps: \(Set(vmaps.all))")
+        #expect(!gradients.all.isEmpty && !vjps.all.isEmpty && !vmaps.all.isEmpty)
+
+        #if DEBUG
+            #expect(EvalLockOwnership.counters.violations == 0)
+        #endif
+    }
+
+    /// A whole training loop per thread: `valueAndGrad(model:)`, an optimizer
+    /// update and `eval` over nested parameter dictionaries.  This is the
+    /// realistic shape of concurrent eval use, and it checks progress (the
+    /// loss must fall) rather than only absence of a crash.
+    ///
+    /// Each thread owns its model, optimizer and data; the inputs are
+    /// deterministic rather than random so the global RNG state -- a separate
+    /// thread-safety question -- is not part of the test.
+    @Test func concurrentTrainingLoops() {
+        final class Model: Module, UnaryLayer {
+            let linear: Linear
+
+            // fixed initial parameters, deliberately wrong, so training has
+            // something to do and the loss trajectory is repeatable
+            init(m: Float, b: Float) {
+                linear = Linear(weight: MLXArray([m], [1, 1]), bias: MLXArray([b]))
+            }
+
+            func callAsFunction(_ x: MLXArray) -> MLXArray { linear(x) }
+        }
+
+        let threads = 4
+        let losses = Collected<(Int, Float, Float)>()
+
+        let finished = onThreads(threads) { t in
+            let model = Model(m: 0, b: 0)
+            eval(model)
+
+            let optimizer = SGD(learningRate: 1e-2)
+            func loss(model: Model, x: MLXArray, y: MLXArray) -> MLXArray {
+                mseLoss(predictions: model(x), targets: y, reduction: .mean)
+            }
+            let lg = valueAndGrad(model: model, loss)
+
+            // a different target per thread, so a leaked gradient or parameter
+            // update from another thread would show up as a stalled loss
+            let m = Float(t + 1)
+            let b = Float(3)
+            let x = MLXArray(Array(stride(from: Float(-2), to: Float(2), by: 0.25)), [16, 1])
+            let y = m * x + b
+
+            var first = Float.nan
+            var last = Float.nan
+            for step in 0 ..< 100 {
+                let (l, grads) = lg(model, x, y)
+                optimizer.update(model: model, gradients: grads)
+                eval(model, optimizer)
+                let value = l.item(Float.self)
+                if step == 0 { first = value }
+                last = value
+            }
+            losses.append((t, first, last))
+        }
+
+        #expect(finished, "a training thread wedged")
+
+        let all = losses.all
+        #expect(all.count == threads)
+        for (t, first, last) in all {
+            #expect(first.isFinite && last.isFinite, "thread \(t): \(first) -> \(last)")
+            #expect(last < first, "thread \(t) did not train: \(first) -> \(last)")
+        }
+    }
+
+    /// The other `evalLock` users -- memory/cache configuration, device
+    /// queries, `tostring` -- run against a stream of evals.  These do not
+    /// wait on the GPU, so all this asserts is that they interleave without
+    /// deadlocking or corrupting results.
+    @Test func concurrentEvalAndOtherEvalLockUsers() {
+        let cacheLimit = Memory.cacheLimit
+        let values = Collected<Float>()
+
+        let finished = onThreads(6) { t in
+            for round in 0 ..< 30 {
+                switch t % 3 {
+                case 0:
+                    let r = heavyGraph(matmuls: 5, size: 128, stream: .default).sum()
+                    eval(r)
+                    values.append(r.item(Float.self))
+                case 1:
+                    _ = Memory.snapshot()
+                    if round % 10 == 0 { Memory.clearCache() }
+                    Memory.cacheLimit = cacheLimit
+                default:
+                    _ = Device.gpu.description
+                    _ = Stream.gpu.description
+                    _ = MLXArray(Float(round)).description
+                }
+            }
+        }
+
+        #expect(finished, "a thread wedged mixing eval with other evalLock users")
+        #expect(values.all.allSatisfy { $0 == Float(128 * 128) }, "values: \(Set(values.all))")
+
+        Memory.cacheLimit = cacheLimit
+    }
+
+    // MARK: Re-entrancy
+
+    /// `evalLock` is recursive on purpose, so an `eval` reached from inside
+    /// something that already holds it must not self-deadlock.  A regression
+    /// here (a non-recursive lock, or a wait performed while holding it in a
+    /// way that cannot complete) hangs, so this is kept small and time
+    /// bounded.
+    @Test func evalIsReentrantWhileHoldingEvalLock() {
+        #if DEBUG
+            #expect(EvalLockOwnership.depth == 0)
+        #endif
+
+        let done = DispatchSemaphore(value: 0)
+        let value = Collected<Float>()
+        let depthInsideEval = Collected<Int>()
+
+        let thread = Thread {
+            withEvalLock {
+                let a = MLXArray(0 ..< 4).asType(.float32) * 2
+                // eval() re-enters withEvalLock, and now also waits while the
+                // outer level is still held
+                eval(a)
+                value.append(a.sum().item(Float.self))
+
+                #if DEBUG
+                    depthInsideEval.append(EvalLockOwnership.depth)
+                #endif
+            }
+            done.signal()
+        }
+        thread.start()
+
+        #expect(
+            done.wait(timeout: .now() + 60) == .success,
+            "eval deadlocked when called with evalLock already held")
+        #expect(value.all == [12])
+
+        #if DEBUG
+            #expect(depthInsideEval.all == [1], "unbalanced evalLock ownership after a nested eval")
+            #expect(EvalLockOwnership.depth == 0)
+        #endif
+    }
+
+    /// `eval` inside a compiled function's traced body: the trace holds
+    /// `evalLock`, so the inner eval is the re-entrant case reached through
+    /// real API rather than through `withEvalLock` directly.
+    ///
+    /// The array evaluated inside the body is already computed and takes no
+    /// part in the trace, so this stays a statement about locking rather than
+    /// about evaluating a traced graph (which MLX does not allow).
+    @Test func evalInsideACompiledBody() {
+        nonisolated(unsafe) var traced = false
+        let side = MLXArray(0 ..< 4).asType(.float32)
+        eval(side)
+
+        let compiled = MLX.compile { (x: MLXArray) in
+            traced = true
+            // re-enters eval (and therefore withEvalLock) while the trace holds it
+            eval(side)
+            return x * 2
+        }
+
+        let r = compiled(MLXArray(Float(3)))
+        eval(r)
+
+        #expect(traced, "the body never traced, so nothing was re-entered")
+        #expect(r.item(Float.self) == 6)
+    }
+
+    // MARK: The reason for the change
+
+    /// `eval` must not hold `evalLock` while waiting for the computation.
+    ///
+    /// Thread A evaluates a long graph on the GPU stream.  Thread B, on the
+    /// CPU stream, does a trivial eval while A's is in flight.  With the wait
+    /// inside the lock B cannot start until A finishes; with the wait outside
+    /// it, B only needs the lock long enough to schedule.
+    ///
+    /// The check is not a bare timing threshold: A publishes a flag for the
+    /// duration of its own `eval` call, and B verifies that flag is *still
+    /// set* at the moment B's own eval returned, which establishes the overlap
+    /// directly.  (B never touches A's arrays -- only the flag.)  Timing is
+    /// used only to tell real serialization ("B took about as long as A") from
+    /// A simply finishing too early to measure, which retries with more work.
+    @Test func evalDoesNotHoldTheLockWhileWaiting() {
+        // Calibrate: how many matmuls are needed for a ~2 second eval?
+        let size = 1024
+        let probeCount = 8
+        let warm = heavyGraph(matmuls: probeCount, size: size, stream: .gpu)
+        eval(warm)
+        let probe = timed {
+            let x = heavyGraph(matmuls: probeCount, size: size, stream: .gpu)
+            eval(x)
+        }
+        let perMatmul = max(probe / Double(probeCount), 1e-5)
+        var matmuls = min(max(Int(2.0 / perMatmul), probeCount), 20_000)
+
+        for attempt in 0 ..< 4 {
+            let heavyIsEvaluating = Flag()
+            let aStarted = DispatchSemaphore(value: 0)
+            let group = DispatchGroup()
+
+            let heavyDurations = Collected<TimeInterval>()
+            let smallDurations = Collected<TimeInterval>()
+            let smallValues = Collected<Float>()
+            let overlapped = Flag()
+
+            let matmulCount = matmuls
+
+            group.enter()
+            let a = Thread {
+                // built on this thread, and never touched by B
+                let heavy = heavyGraph(matmuls: matmulCount, size: size, stream: .gpu)
+                aStarted.signal()
+                heavyIsEvaluating.set(true)
+                let duration = timed { eval(heavy) }
+                heavyIsEvaluating.set(false)
+                heavyDurations.append(duration)
+                group.leave()
+            }
+
+            group.enter()
+            let b = Thread {
+                guard aStarted.wait(timeout: .now() + 60) == .success else {
+                    group.leave()
+                    return
+                }
+                // let A get past scheduling and into its wait
+                Thread.sleep(forTimeInterval: 0.25)
+
+                let cpu = StreamOrDevice.cpu
+                let x = MLXArray(Float(2))
+                let r = MLX.multiply(x, MLXArray(Float(21)), stream: cpu)
+                let duration = timed { eval(r) }
+                // read immediately: if A is still inside its eval, B's eval
+                // completed while A was waiting
+                overlapped.set(heavyIsEvaluating.isSet)
+                smallDurations.append(duration)
+                smallValues.append(r.item(Float.self))
+                group.leave()
+            }
+
+            a.start()
+            b.start()
+
+            #expect(group.wait(timeout: .now() + 600) == .success, "a thread wedged")
+            #expect(smallValues.all == [42])
+
+            guard let heavyDuration = heavyDurations.all.first,
+                let smallDuration = smallDurations.all.first
+            else {
+                Issue.record("a thread did not report a duration")
+                return
+            }
+
+            if overlapped.isSet {
+                // B ran to completion while A's eval was still in flight
+                #expect(
+                    smallDuration < heavyDuration,
+                    """
+                    the small eval took \(smallDuration)s while the long eval took \
+                    \(heavyDuration)s
+                    """)
+                return
+            }
+
+            // A finished before B could look.  Serialization and "A was simply
+            // too fast" are distinguished by how long B waited.
+            if smallDuration >= heavyDuration / 2 {
+                Issue.record(
+                    """
+                    a trivial eval on another stream took \(smallDuration)s while a \
+                    \(heavyDuration)s eval was in flight -- eval appears to hold evalLock while \
+                    waiting
+                    """)
+                return
+            }
+
+            // grow the window and try again
+            matmuls = min(matmuls * 4, 40_000)
+            print(
+                """
+                evalDoesNotHoldTheLockWhileWaiting: attempt \(attempt) inconclusive \
+                (long eval \(heavyDuration)s, small eval \(smallDuration)s); retrying with \
+                \(matmuls) matmuls
+                """)
+        }
+
+        Issue.record(
+            """
+            unable to keep a long eval in flight long enough to observe overlap -- \
+            the concurrency of eval was not established by this run
+            """)
+    }
+}
diff --git a/skills/mlx-swift/references/concurrency.md b/skills/mlx-swift/references/concurrency.md
index c4d4fc7c8..96e4720a6 100644
--- a/skills/mlx-swift/references/concurrency.md
+++ b/skills/mlx-swift/references/concurrency.md
@@ -44,26 +44,53 @@ Task {
 
 ## evalLock
 
-Evaluation and stream creation are serialized through a global `evalLock`:
-
-```swift
-// Defined in Source/MLX/Transforms+Eval.swift
-let evalLock = NSRecursiveLock()
-
-// Used in eval():
-func eval(_ arrays: MLXArray...) {
-    _ = evalLock.withLock {
-        mlx_eval(...)
-    }
-}
-```
+Evaluation, function transformations (`grad`, `valueAndGrad`, `vjp`, `jvp`,
+`vmap`, `compile`, export) and stream creation are serialized through a global,
+recursive `evalLock`.
 
 This means:
-- `eval()` calls and stream creation are serialized across threads
+- **`eval()` holds `evalLock` only while scheduling the work**, not while
+  waiting for it to finish. `eval()` is still synchronous -- the arrays are
+  computed when it returns -- but evaluations on different threads overlap
+  their GPU execution: a small eval on one thread does not wait for a large
+  eval on another thread to complete, only for it to be scheduled.
+- Inside a transformation (or any other code already holding `evalLock`),
+  `eval()` is fully synchronous under the lock.
 - **Important**: evalLock only protects eval/stream operations, NOT all GPU ops
 - Lazy array operations are NOT automatically thread-safe
 - Parallelism happens within Metal, not between MLX calls
 
+### Evaluating inside transformations
+
+- Inside `grad` / `valueAndGrad` / `vjp` / `jvp` (including custom VJPs), a
+  synchronous `eval()` of an intermediate value is allowed, as is reading it
+  (`item()`, `asArray()`, control flow on a value).  `asyncEval()` is **not**
+  allowed there.
+- Inside `compile` and `vmap`, evaluating anything derived from the inputs is
+  not allowed with either `eval()` or `asyncEval()` -- these trace with
+  placeholder inputs that have no data.
+
+```swift
+let g = grad { (x: MLXArray) -> MLXArray in
+    let y = x * x
+    print(y.asArray(Float.self))  // OK: synchronous eval of an intermediate
+    // asyncEval(y)               // error: not allowed inside a graph transformation
+    return y.sum()
+}
+```
+
+### Arrays shared across threads
+
+`MLXArray` is not thread safe, and that includes evaluation.  If several
+threads evaluate graphs that contain the *same* unevaluated array (e.g. lazily
+loaded model weights shared by several inference threads), they race on that
+array's state.  Evaluate shared arrays once, on one thread, before sharing them:
+
+```swift
+let model = try loadModel(...)
+eval(model)   // realize the weights before handing the model to other threads
+```
+
 ## @unchecked Sendable Types
 
 These types are marked `@unchecked Sendable` and are safe to share: