Sync upstream ml-explore/mlx-swift main (2f96983)

.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.

.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

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

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"

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

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)

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).

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,18 +182,22 @@ 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)
 
         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)
         }
     }
 

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,48 +104,107 @@ 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`
 ///
 /// ### 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.