Sync upstream ml-explore/mlx-swift main (43c90af)

.github/workflows/release.yml

@@ -45,7 +45,7 @@ jobs:
         shell: bash
         working-directory: xcode/build/output
         run: |
-          for framework in MLX Cmlx cLogSupport MLXNN MLXOptimizers; do
+          for framework in MLX Cmlx MLXNN MLXOptimizers; do
             zip -r -y "${framework}.xcframework.zip" "${framework}.xcframework"
           done
 
@@ -66,6 +66,5 @@ 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

CMakeLists.txt

@@ -100,20 +100,6 @@ if(MLX_BUILD_CUDA)
   )
 endif()
 
-# cLogSupport package small C shim that exposes os_log_type_enabled() to Swift
-# (see #491)
-set(cLogSupport_include ${CMAKE_CURRENT_LIST_DIR}/Source/cLogSupport/include)
-set(cLogSupport_module_dir ${CMAKE_CURRENT_BINARY_DIR}/cLogSupport-module)
-file(
-  WRITE ${cLogSupport_module_dir}/module.modulemap
-  "module cLogSupport {\n    header \"${cLogSupport_include}/cLogSupport.h\"\n    export *\n}\n"
-)
-
-add_library(cLogSupport STATIC
-            ${CMAKE_CURRENT_LIST_DIR}/Source/cLogSupport/cLogSupport.c)
-target_include_directories(cLogSupport PUBLIC ${cLogSupport_include}
-                                              ${cLogSupport_module_dir})
-
 file(GLOB MLX-src ${CMAKE_CURRENT_LIST_DIR}/Source/MLX/*.swift)
 
 if(NOT MLX_BUILD_METAL)
@@ -129,7 +115,7 @@ add_library(MLX STATIC ${MLX-src})
 target_include_directories(MLX
                            PUBLIC ${CMAKE_CURRENT_LIST_DIR}/Source/Cmlx/include)
 # Note: mlxc is target name exported by mlx-c
-target_link_libraries(MLX PUBLIC mlxc Numerics cLogSupport)
+target_link_libraries(MLX PUBLIC mlxc Numerics)
 target_compile_options(MLX PRIVATE -package-name MLX)
 
 # MLXRandom package

Package.swift

@@ -281,17 +281,10 @@ let package = Package(
             dependencies: ["Cmlx"]
         ),
 
-        .target(
-            name: "cLogSupport",
-            path: "Source/cLogSupport",
-            publicHeadersPath: "include"
-        ),
-
         .target(
             name: "MLX",
             dependencies: [
                 "Cmlx",
-                "cLogSupport",
                 .product(name: "Numerics", package: "swift-numerics"),
             ],
             exclude: mlxSwiftExcludes,

Source/MLX/Documentation.docc/Articles/compilation.md

@@ -263,8 +263,8 @@ class ScaleModule: Module, UnaryLayer {
 
 let model = ScaleModule(2.0)
 
-// no inputs/outputs -- the module's parameters are captured as constants
-let compiled = compile { (x: MLXArray) -> MLXArray in
+// `inputs: []` -- the module's parameters are captured as constants
+let compiled = compile(inputs: []) { (x: MLXArray) -> MLXArray in
     model(x)
 }
 
@@ -278,6 +278,12 @@ eval(model)
 print(compiled(MLXArray(Float(3))))
 ```
 
+Note that this pitfall is hard to hit by accident: a plain `compile { ... }`
+closure must be `@Sendable` (see <doc:#Sendable> below), and capturing a
+`Module` (which is not `Sendable`) is a compile error. Writing
+`compile(inputs: []) { ... }` is the explicit way to say "capture this state
+and freeze it at trace time".
+
 To make the compiled function observe the update, pass the module via `inputs:`:
 
 ```swift
@@ -302,6 +308,81 @@ graph into the real arrays, and is only needed when the compiled closure itself
 below). Since `model(x)` above only reads `scale`, `inputs: [model]` alone is
 enough; adding `outputs: [model]` would be harmless but unnecessary.
 
+## Sendable
+
+There are two families of `compile` functions, and which one is used determines
+whether the result is `@Sendable`:
+
+- ``compile(shapeless:_:)-(Bool,([MLXArray])->[MLXArray])`` (and variants) takes
+  a `@Sendable` function and returns a `@Sendable` function. A closure that
+  captures nothing -- it only uses its arguments -- is `@Sendable`, so this is
+  what you get for the common case of compiling a pure function:
+
+    ```swift
+    // usable as a global or across concurrency domains
+    let compiledGelu: @Sendable (MLXArray) -> MLXArray = compile(shapeless: true) { x in
+        x * (1 + erf(x / sqrt(2))) / 2
+    }
+    ```
+
+- ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])``
+  (and variants) takes any function, including one that captures non-`Sendable`
+  state such as an ``MLXArray`` or `Module`, and returns a function that is
+  **not** `@Sendable`. It is selected whenever you pass `inputs:` or `outputs:`,
+  or when the function you pass is not `@Sendable`.
+
+If a `compile { ... }` closure fails to build with an error like
+"capture of 'model' with non-Sendable type ... in a '@Sendable' closure", the
+closure is capturing state. Either:
+
+- list the captured state in `inputs:` (and `outputs:` if the closure mutates
+  it) so the compiled function sees updates, or
+- write `compile(inputs: []) { ... }` to deliberately capture the values as
+  constants at trace time (see the `Module` example above). This is also the
+  spelling to use for captured values that are not ``Updatable`` at all.
+
+This split matches how the result is used: a compiled function that captures
+mutable, non-`Sendable` state cannot safely be shared between threads, while a
+compiled pure function can.
+
+### Materialized state
+
+``MaterializedArray`` and `MaterializedModule` (from `MLXNN`) are `Sendable`, so
+a `@Sendable` closure may capture them and still use the `@Sendable` family:
+
+```swift
+let model = MaterializedModule(Linear(16, 16))
+
+// @Sendable -- the model's parameters are captured as constants
+let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in
+    model(x)
+}
+```
+
+Captured values are frozen at trace time, which is exactly right here: the
+contents of a ``MaterializedArray`` (and the parameters of a `MaterializedModule`)
+are sealed and can never change, so there is nothing to observe and no need for
+`inputs:`.
+
+What *can* change is which value the closure reaches.  If it captures a
+mutable container -- for example a `Mutex<MaterializedModule>` or a class that
+holds the current model -- and that container is later pointed at a different
+model, the compiled function keeps using the one it was traced with.  Compile a
+new function when the model is replaced.
+
+Sealed state reached through `inputs:` is also treated as a constant rather
+than being fed in on every call, but there the compiled function checks the
+identity of the sealed arrays on each call and retraces if they have been
+replaced.
+
+> Note: in the Swift 5 language mode with minimal concurrency checking the
+compiler does not diagnose non-`Sendable` captures, so a capturing closure
+literal may still resolve to the `@Sendable` variant. Use the Swift 6 language
+mode (or `-strict-concurrency=complete`) to get these checks.
+
+See also ``vmapSendable(_:inAxes:outAxes:)-3tm2r`` for the equivalent in
+<doc:vmap>.
+
 ## Compiling Training Graphs 
 
 This section will step through how to use ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` 
@@ -384,4 +465,7 @@ for _ in 0 ..< 30 {
 - ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])``
 - ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,(MLXArray)->MLXArray)``
 - ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,(MLXArray,MLXArray)->MLXArray)``
+- ``compile(shapeless:_:)-(Bool,([MLXArray])->[MLXArray])``
+- ``compile(shapeless:_:)-(Bool,(MLXArray)->MLXArray)``
+- ``compile(shapeless:_:)-(Bool,(MLXArray,MLXArray)->MLXArray)``
 - ``compile(enable:)``

Source/MLX/Documentation.docc/Articles/vmap.md

@@ -44,3 +44,15 @@ Here `x` is mapped over its first axis while `y` is used as a broadcast value.
 You can nest calls to `vmap` to map over multiple axes.
 Each nested `vmap` introduces another batch dimension in the result.
 
+## Sendable
+
+If the function being transformed is `Sendable` (typically meaning it
+does not capture any MLXArray values, it only uses its inputs) you can use
+``vmapSendable(_:inAxes:outAxes:)-3tm2r`` (and variants) to
+produce a `Sendable` result closure:
+
+```swift
+@Sendable
+func add(_ x: MLXArray, _ y: MLXArray) -> MLXArray { x + y }
+let vf = vmapSendable(add, inAxes: (0, nil))
+```

Source/MLX/Documentation.docc/MLX.md

@@ -92,6 +92,15 @@ See <doc:wired-memory> for configuration, best practices, and policy guidance.
 
 - ``MLXArray``
 
+### Concurrency
+
+`MLXArray` is not `Sendable`.  To share an array or model across task and actor
+boundaries, take an evaluated, immutable snapshot.
+
+- ``MaterializedArray``
+- ``materialize(_:)->MaterializedArray``
+- ``materialize(_:)->[MaterializedArray]``
+
 ### Free Functions
 
 - <doc:free-functions>

Source/MLX/Documentation.docc/MLXArray.md

@@ -42,6 +42,11 @@ let c = a + b
 may be large unresolved graphs) and the `addition` operator that combines them.  It is not safe
 to create `c` in one thread and consume/evaluate it in another.
 
+If you need to share an array across a task or actor boundary, take a
+``MaterializedArray`` snapshot via ``MLXArray/materialized()`` (or
+``materialize(_:)->MaterializedArray``).  A `MaterializedArray` is fully
+evaluated, immutable, and `Sendable`.
+
 ## Memory Safety
 
 > `MLXArray` is not memory safe.
@@ -91,3 +96,10 @@ and ``shape`` of the dimensions without changing the number of elements or conte
 ### Conversion and Data Types
 
 - <doc:conversion>
+
+### Concurrency
+
+- ``materialized()``
+- ``MaterializedArray``
+- ``materialize(_:)->MaterializedArray``
+- ``materialize(_:)->[MaterializedArray]``

Source/MLX/Documentation.docc/free-functions.md

@@ -219,6 +219,11 @@ operations as methods for convenience.
 - ``jvp(_:primals:tangents:)``
 - ``vjp(_:primals:cotangents:)``
 
+### Concurrency
+
+- ``materialize(_:)->MaterializedArray``
+- ``materialize(_:)->[MaterializedArray]``
+
 ### Device
 
 - ``using(device:fn:)``

Source/MLX/GPU+Metal.swift

@@ -242,13 +242,29 @@ public enum GPU {
     }
 
     /// Get information about the GPU device and system settings
+    ///
+    /// The value is computed once per process.
     public static func deviceInfo() -> DeviceInfo {
+        cachedDeviceInfo
+    }
+
+    private static let cachedDeviceInfo: DeviceInfo = {
         var mib = [CTL_HW, HW_MEMSIZE]
         var memSize: size_t = 0
         var length: size_t = MemoryLayout.size(ofValue: memSize)
         sysctl(&mib, 2, &memSize, &length, nil, 0)
 
-        if let device = MTLCreateSystemDefaultDevice() {
+        // Select the device the same way as MLX (metal::load_device): the first
+        // device of MTLCopyAllDevices(), else the system default device. In a
+        // command-line process MTLCreateSystemDefaultDevice() opens a window
+        // server session, which can stall for several seconds.
+        #if os(macOS)
+            let device = MTLCopyAllDevices().first ?? MTLCreateSystemDefaultDevice()
+        #else
+            let device = MTLCreateSystemDefaultDevice()
+        #endif
+
+        if let device {
             let architecture: String
             if #available(macOS 14.0, iOS 17.0, tvOS 17.0, *) {
                 architecture = device.architecture.name
@@ -265,5 +281,5 @@ public enum GPU {
                 architecture: "Unknown", maxBufferSize: 0, maxRecommendedWorkingSetSize: 0,
                 memorySize: memSize)
         }
-    }
+    }()
 }

Source/MLX/Logging.swift

@@ -223,23 +223,23 @@ extension MLXLogger {
 
 #if canImport(OSLog)
     import OSLog
-    import cLogSupport
 
     /// On systems that support it, a log handler that is backed by OSLog.
     public struct OSLogHandler: MLXLogHandler {
 
+        // Logger.isEnabled(type:) is absent on iOS 26.2.
+        // This handle checks the level before message evaluation.
+        private let _osLog: OSLog
         private let logger: os.Logger
 
-        // see #491
-        private let osLog: OSLog
-
         public static func install() {
             MLXLogger.factory = { OSLogHandler(label: $0) }
         }
 
         public init(label: String) {
-            self.logger = Logger(subsystem: "mlx-swift", category: label)
-            self.osLog = OSLog(subsystem: "mlx-swift", category: label)
+            let osLog = OSLog(subsystem: "mlx-swift", category: label)
+            self._osLog = osLog
+            self.logger = Logger(osLog)
         }
 
         public func log(
@@ -255,7 +255,7 @@ extension MLXLogger {
                 case .warning: .error
                 case .error: .error
                 }
-            if is_os_log_enabled(osLog, level) {
+            if _osLog.isEnabled(type: level) {
                 let message = message()
                 logger.log(level: level, "\(message, privacy: .public)")
             }

Source/MLX/MLXArray+Indexing.swift

@@ -388,8 +388,7 @@ extension MLXArray {
                 var result = mlx_array_new()
                 mlx_scatter(
                     &result, self.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx)
-                mlx_array_set(&self.ctx, result)
-                mlx_array_free(result)
+                self._updateInternal(MLXArray(result))
                 return
             } else {
                 self._updateInternal(update)

Source/MLX/MLXArray+Init.swift

@@ -683,39 +683,3 @@ extension MLXArray {
     }
 
 }
-
-// MARK: - Expressible by literals
-
-extension MLXArray: ExpressibleByArrayLiteral {
-
-    // Note: MLXArray does not implement ExpressibleByFloatLiteral etc. because
-    // we want to create arrays in the context of the other arrays.  For example:
-    //
-    // let x = MLXArray(1.5, dtype: .float16)
-    // let r = x + 2.5
-    //
-    // We expect r to have a dtype of float16.  See ``ScalarOrArray``.
-
-    /// Initializer allowing creation of 1d `MLXArray` from an array literal.
-    ///
-    /// ```swift
-    /// let a: MLXArray = [1, 2, 3]
-    /// ```
-    ///
-    /// This is convenient for methods that have `MLXArray` parameters:
-    ///
-    /// ```swift
-    /// print(array.take([1, 2, 3], axis: 0))
-    /// ```
-    ///
-    /// ### See Also
-    /// - <doc:initialization>
-    public convenience init(arrayLiteral elements: Int32...) {
-        let ctx = elements.withUnsafeBufferPointer { ptr in
-            let shape = [Int32(elements.count)]
-            return mlx_array_new_data(
-                ptr.baseAddress!, shape, Int32(shape.count), Int32.dtype.cmlxDtype)
-        }
-        self.init(ctx)
-    }
-}