🔄 Auto-Sync: ml-explore/mlx-swift main (863ee9f, 4 commit(s))

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,51 @@ 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.
+
+> 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 +435,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/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/MLXCustomFunction.swift

@@ -1,26 +1,39 @@
 import Cmlx
 import Foundation
 
-public enum MLXCustomFunctionComponent {
-    case forward(([MLXArray]) -> [MLXArray])
-    case vjp(([MLXArray], [MLXArray]) -> [MLXArray])
+public enum MLXCustomFunctionComponent: Sendable {
+    case forward(@Sendable ([MLXArray]) -> [MLXArray])
+    case vjp(@Sendable ([MLXArray], [MLXArray], [MLXArray]) -> [MLXArray])
 }
 
-public func Forward(_ f: @escaping ([MLXArray]) -> [MLXArray]) -> MLXCustomFunctionComponent {
+public func Forward(_ f: @escaping @Sendable ([MLXArray]) -> [MLXArray])
+    -> MLXCustomFunctionComponent
+{
     .forward(f)
 }
 
-public func VJP(_ f: @escaping ([MLXArray], [MLXArray]) -> [MLXArray]) -> MLXCustomFunctionComponent
-{
+public func VJP(
+    _ f: @escaping @Sendable (_ primals: [MLXArray], _ cotangents: [MLXArray]) -> [MLXArray]
+) -> MLXCustomFunctionComponent {
+    .vjp { primals, cotangents, _ in
+        f(primals, cotangents)
+    }
+}
+
+public func VJP(
+    _ f:
+        @escaping @Sendable (_ primals: [MLXArray], _ cotangents: [MLXArray], _ outputs: [MLXArray])
+        -> [MLXArray]
+) -> MLXCustomFunctionComponent {
     .vjp(f)
 }
 
 final class _CustomFunctionState: @unchecked Sendable {
 
     private let lock = NSLock()
 
-    private let forwardFn: ([MLXArray]) -> [MLXArray]
-    private let vjpFn: (([MLXArray], [MLXArray]) -> [MLXArray])?
+    private let forwardFn: @Sendable ([MLXArray]) -> [MLXArray]
+    private let vjpFn: (@Sendable ([MLXArray], [MLXArray], [MLXArray]) -> [MLXArray])?
 
     private var forwardClosure: mlx_closure!
     private var vjpClosure: mlx_closure_custom!
@@ -29,8 +42,8 @@ final class _CustomFunctionState: @unchecked Sendable {
     private var combined: mlx_closure!
 
     init(
-        forward: @escaping ([MLXArray]) -> [MLXArray],
-        vjp: (([MLXArray], [MLXArray]) -> [MLXArray])?
+        forward: @escaping @Sendable ([MLXArray]) -> [MLXArray],
+        vjp: (@Sendable ([MLXArray], [MLXArray], [MLXArray]) -> [MLXArray])?
     ) {
         self.forwardFn = forward
         self.vjpFn = vjp
@@ -53,7 +66,7 @@ final class _CustomFunctionState: @unchecked Sendable {
                     Unmanaged<AnyObject>
                     .fromOpaque(payload!)
                     .takeUnretainedValue()
-                    as! ([MLXArray]) -> [MLXArray]
+                    as! @Sendable ([MLXArray]) -> [MLXArray]
 
                 let inp = mlx_vector_array_values(inputs)
                 let result = swiftFn(inp)
@@ -67,16 +80,17 @@ final class _CustomFunctionState: @unchecked Sendable {
 
         if let vjpFn = vjpFn {
             vjpClosure = mlx_closure_custom_new_func_payload(
-                { out, primals, cotangents, _, payload in
+                { out, primals, cotangents, outputs, payload in
                     let fn =
                         Unmanaged<AnyObject>
                         .fromOpaque(payload!)
                         .takeUnretainedValue()
-                        as! ([MLXArray], [MLXArray]) -> [MLXArray]
+                        as! @Sendable ([MLXArray], [MLXArray], [MLXArray]) -> [MLXArray]
 
                     let p = mlx_vector_array_values(primals)
                     let c = mlx_vector_array_values(cotangents)
-                    out!.pointee = new_mlx_vector_array(fn(p, c))
+                    let o = mlx_vector_array_values(outputs)
+                    out!.pointee = new_mlx_vector_array(fn(p, c, o))
                     return 0
                 },
                 Unmanaged.passRetained(vjpFn as AnyObject).toOpaque()
@@ -122,10 +136,10 @@ public enum MLXCustomFunctionBuilder {
 
         public static func buildBlock(
             _ components: MLXCustomFunctionComponent...
-        ) -> ([MLXArray]) -> [MLXArray] {
+        ) -> @Sendable ([MLXArray]) -> [MLXArray] {
 
-            var forwardFn: (([MLXArray]) -> [MLXArray])?
-            var vjpFn: (([MLXArray], [MLXArray]) -> [MLXArray])?
+            var forwardFn: (@Sendable ([MLXArray]) -> [MLXArray])?
+            var vjpFn: (@Sendable ([MLXArray], [MLXArray], [MLXArray]) -> [MLXArray])?
 
             for c in components {
                 switch c {
@@ -146,7 +160,7 @@ public enum MLXCustomFunctionBuilder {
 }
 
 public func CustomFunction(
-    @MLXCustomFunctionBuilder.Builder _ build: () -> ([MLXArray]) -> [MLXArray]
-) -> ([MLXArray]) -> [MLXArray] {
+    @MLXCustomFunctionBuilder.Builder _ build: () -> @Sendable ([MLXArray]) -> [MLXArray]
+) -> @Sendable ([MLXArray]) -> [MLXArray] {
     build()
 }