From 26d685a284ecbd7c30e9c2aec7e923ab21c4587d Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Mon, 19 Oct 2026 08:55:16 -0700
Subject: [PATCH 1/4] compile and vmap can produce Sendable closures from
 Sendable inputs (#469)

* compile and vmap can produce Sendable closures from Sendable inputs

- this idea comes from https://github.com/ml-explore/mlx-swift-lm/pull/589 (@aleroot)
- breaking change for compile calls that capture state without using inputs:
---
 .../Articles/compilation.md                   |  58 +-
 .../MLX/Documentation.docc/Articles/vmap.md   |  12 +
 Source/MLX/Transforms+Compile.swift           | 118 +++-
 Source/MLX/Transforms+CompileOverloads.swift  | 662 ++++++++++++++++--
 Source/MLX/Transforms+Grad.swift              |   4 +-
 Source/MLX/Transforms+Vmap.swift              | 118 +++-
 Tests/MLXTests/CompileLockOrderingTests.swift |   7 +-
 Tests/MLXTests/EvalTests.swift                |  26 -
 Tests/MLXTests/OptimizerTests.swift           |   1 +
 Tests/MLXTests/TransformTests.swift           |  38 +-
 Tests/MLXTests/WiredMemoryTests.swift         |   2 +-
 11 files changed, 938 insertions(+), 108 deletions(-)

diff --git a/Source/MLX/Documentation.docc/Articles/compilation.md b/Source/MLX/Documentation.docc/Articles/compilation.md
index 736851f29..b2f87e295 100644
--- a/Source/MLX/Documentation.docc/Articles/compilation.md
+++ b/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:)``
diff --git a/Source/MLX/Documentation.docc/Articles/vmap.md b/Source/MLX/Documentation.docc/Articles/vmap.md
index 719e631f3..6a8887242 100644
--- a/Source/MLX/Documentation.docc/Articles/vmap.md
+++ b/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))
+```
diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift
index c9616f753..527a82e44 100644
--- a/Source/MLX/Transforms+Compile.swift
+++ b/Source/MLX/Transforms+Compile.swift
@@ -3,7 +3,10 @@
 import Cmlx
 import Foundation
 
-// Note: this is all immutable state -- the `id` property is only set at init time
+// Note: this is all immutable state -- the `id` property is only set at init time.
+// Also: this is marked as Sendable even though it may capture non-Sendable state
+// in inputs/outputs.  The functions that create it will not return a Sendable
+// type in that case -- see compile(inputs:)
 final class CompiledFunction: @unchecked (Sendable) {
 
     /// unique (for the lifetime of the object) identifier for the compiled function
@@ -38,6 +41,17 @@ final class CompiledFunction: @unchecked (Sendable) {
         self.id = UInt(bitPattern: Unmanaged.passUnretained(self).toOpaque())
     }
 
+    init(
+        shapeless: Bool,
+        _ f: @escaping @Sendable ([MLXArray]) -> [MLXArray]
+    ) {
+        self.f = f
+        self.inputs = []
+        self.outputs = []
+        self.shapeless = shapeless
+        self.id = UInt(bitPattern: Unmanaged.passUnretained(self).toOpaque())
+    }
+
     /// Remove the compiled structure from the back end.
     ///
     /// ### Why this takes `evalLock`
@@ -287,7 +301,7 @@ final class CompiledFunction: @unchecked (Sendable) {
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping ([MLXArray]) -> [MLXArray]
-) -> @Sendable ([MLXArray]) -> [MLXArray] {
+) -> ([MLXArray]) -> [MLXArray] {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless, f)
 
     return { arrays in
@@ -295,6 +309,31 @@ public func compile(
     }
 }
 
+/// Returns a compiled function that produces the same output as `f()`.
+///
+/// - Parameters:
+///   - shapeless: A function compiled with the `shapeless`
+///     option enabled will not be recompiled when the input shape changes. Not all
+///     functions can be compiled with `shapeless` enabled. Attempting to compile
+///     such functions with shapeless enabled will throw. Note, changing the number
+///     of dimensions or type of any input will result in a recompilation even with
+///     `shapeless` set to `true`
+///   - f: function to compile
+/// - Returns: a new function that produces the same output as `f()`
+///
+/// ### See Also
+/// - <doc:compilation>
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable ([MLXArray]) -> [MLXArray]
+) -> @Sendable ([MLXArray]) -> [MLXArray] {
+    let compileState = CompiledFunction(shapeless: shapeless, f)
+
+    return { arrays in
+        compileState.call(arrays)
+    }
+}
+
 /// Overload of ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` that takes a single ``MLXArray`` and
 /// produces a single ``MLXArray``.
 ///
@@ -304,7 +343,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray) -> MLXArray
-) -> @Sendable (MLXArray) -> MLXArray {
+) -> (MLXArray) -> MLXArray {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         [f($0[0])]
     }
@@ -318,6 +357,29 @@ public func compile(
     }
 }
 
+/// Overload of ``compile(shapeless:_:)-6mb2r`` that takes a single ``MLXArray`` and
+/// produces a single ``MLXArray``.
+///
+/// ### See Also
+/// - <doc:compilation>
+/// - ``compile(inputs:outputs:shapeless:_:)-7korq``
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray) -> MLXArray
+) -> @Sendable (MLXArray) -> MLXArray {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        [f($0[0])]
+    }
+
+    return { a in
+        let r = compileState.call([a])
+        // r is empty only when an MLX error fired inside a withError scope — the
+        // error is already stored in the ErrorBox.  Return a placeholder so that
+        // withError can throw instead of crashing with "Index out of range".
+        return r.isEmpty ? MLXArray(0) : r[0]
+    }
+}
+
 /// Overload of ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` that takes two ``MLXArray`` and
 /// produces a single ``MLXArray``.
 ///
@@ -328,7 +390,7 @@ public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray) -> MLXArray
 )
-    -> @Sendable (MLXArray, MLXArray) -> MLXArray
+    -> (MLXArray, MLXArray) -> MLXArray
 {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         [f($0[0], $0[1])]
@@ -340,6 +402,28 @@ public func compile(
     }
 }
 
+/// Overload of ``compile(shapeless:_:)-6mb2r`` that takes two ``MLXArray`` and
+/// produces a single ``MLXArray``.
+///
+/// ### See Also
+/// - <doc:compilation>
+/// - ``compile(inputs:outputs:shapeless:_:)-7korq``
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray) -> MLXArray
+)
+    -> @Sendable (MLXArray, MLXArray) -> MLXArray
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        [f($0[0], $0[1])]
+    }
+
+    return { a, b in
+        let r = compileState.call([a, b])
+        return r.isEmpty ? MLXArray(0) : r[0]
+    }
+}
+
 /// Overload of ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` that takes three ``MLXArray`` and
 /// produces a single ``MLXArray``.
 ///
@@ -348,9 +432,9 @@ public func compile(
 /// - ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])``
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
-    _ f: @Sendable @escaping (MLXArray, MLXArray, MLXArray) -> MLXArray
+    _ f: @escaping (MLXArray, MLXArray, MLXArray) -> MLXArray
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray) -> MLXArray
+    -> (MLXArray, MLXArray, MLXArray) -> MLXArray
 {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         [f($0[0], $0[1], $0[2])]
@@ -362,6 +446,28 @@ public func compile(
     }
 }
 
+/// Overload of ``compile(shapeless:_:)-6mb2r`` that takes three ``MLXArray`` and
+/// produces a single ``MLXArray``.
+///
+/// ### See Also
+/// - <doc:compilation>
+/// - ``compile(inputs:outputs:shapeless:_:)-7korq``
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray, MLXArray) -> MLXArray
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray) -> MLXArray
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        [f($0[0], $0[1], $0[2])]
+    }
+
+    return { a, b, c in
+        let r = compileState.call([a, b, c])
+        return r.isEmpty ? MLXArray(0) : r[0]
+    }
+}
+
 /// Globally enable or disable ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])``.
 ///
 /// Default is enabled.
diff --git a/Source/MLX/Transforms+CompileOverloads.swift b/Source/MLX/Transforms+CompileOverloads.swift
index c10978462..2f478dfb5 100644
--- a/Source/MLX/Transforms+CompileOverloads.swift
+++ b/Source/MLX/Transforms+CompileOverloads.swift
@@ -2,11 +2,13 @@
 
 import Cmlx
 
+// MARK: - inputs:/outputs:/not Sendable
+
 @_documentation(visibility: internal)
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray) -> (MLXArray, MLXArray)
-) -> @Sendable (MLXArray) -> (MLXArray, MLXArray) {
+) -> (MLXArray) -> (MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0])
@@ -22,7 +24,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray) -> (MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+) -> (MLXArray) -> (MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0])
@@ -38,7 +40,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
+) -> (MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0])
@@ -54,7 +56,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray) -> (MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+) -> (MLXArray, MLXArray) -> (MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1])
@@ -70,7 +72,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+) -> (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1])
@@ -86,7 +88,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
+) -> (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1])
@@ -102,7 +104,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+) -> (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2])
@@ -118,7 +120,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+) -> (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2])
@@ -133,8 +135,11 @@ public func compile(
 @_documentation(visibility: internal)
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
-    _ f: @escaping (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
+    _ f:
+        @escaping (MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray, MLXArray
+        )
+) -> (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2])
@@ -150,7 +155,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+) -> (MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         [f(args[0], args[1], args[2], args[3])]
@@ -164,7 +169,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+) -> (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2], args[3])
@@ -179,8 +184,11 @@ public func compile(
 @_documentation(visibility: internal)
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
-    _ f: @escaping (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+    _ f:
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray
+        )
+) -> (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2], args[3])
@@ -199,8 +207,7 @@ public func compile(
         @escaping (MLXArray, MLXArray, MLXArray, MLXArray) -> (
             MLXArray, MLXArray, MLXArray, MLXArray
         )
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
-{
+) -> (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2], args[3])
@@ -216,7 +223,7 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f: @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+) -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         [f(args[0], args[1], args[2], args[3], args[4])]
@@ -229,8 +236,11 @@ public func compile(
 @_documentation(visibility: internal)
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
-    _ f: @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+    _ f:
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray
+        )
+) -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2], args[3], args[4])
@@ -249,8 +259,7 @@ public func compile(
         @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
             MLXArray, MLXArray, MLXArray
         )
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
-{
+) -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2], args[3], args[4])
@@ -270,7 +279,7 @@ public func compile(
             MLXArray, MLXArray, MLXArray, MLXArray
         )
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
         MLXArray, MLXArray, MLXArray, MLXArray
     )
 {
@@ -288,8 +297,9 @@ public func compile(
 @_documentation(visibility: internal)
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
-    _ f: @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+    _ f:
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
+) -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         [f(args[0], args[1], args[2], args[3], args[4], args[5])]
@@ -306,8 +316,7 @@ public func compile(
         @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
             MLXArray, MLXArray
         )
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
-{
+) -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         let r = f(args[0], args[1], args[2], args[3], args[4], args[5])
@@ -327,7 +336,7 @@ public func compile(
             MLXArray, MLXArray, MLXArray
         )
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
         MLXArray, MLXArray, MLXArray
     )
 {
@@ -350,7 +359,7 @@ public func compile(
             MLXArray, MLXArray, MLXArray, MLXArray
         )
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
         MLXArray, MLXArray, MLXArray, MLXArray
     )
 {
@@ -369,8 +378,9 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
-) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> MLXArray
+) -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
         args in
         [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])]
@@ -384,11 +394,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> (
             MLXArray, MLXArray
         )
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
         MLXArray, MLXArray
     )
 {
@@ -407,11 +418,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> (
             MLXArray, MLXArray, MLXArray
         )
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
         MLXArray, MLXArray, MLXArray
     )
 {
@@ -430,11 +442,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> (
             MLXArray, MLXArray, MLXArray, MLXArray
         )
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
         MLXArray, MLXArray, MLXArray, MLXArray
     )
 {
@@ -453,10 +466,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        @escaping (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
         -> MLXArray
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
     MLXArray
 {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
@@ -472,10 +487,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        @escaping (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
         -> (MLXArray, MLXArray)
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
     (MLXArray, MLXArray)
 {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
@@ -493,10 +510,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        @escaping (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
         -> (MLXArray, MLXArray, MLXArray)
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
     (MLXArray, MLXArray, MLXArray)
 {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
@@ -514,10 +533,12 @@ public func compile(
 public func compile(
     inputs: [any Updatable] = [], outputs: [any Updatable] = [], shapeless: Bool = false,
     _ f:
-        @escaping (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        @escaping (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
         -> (MLXArray, MLXArray, MLXArray, MLXArray)
 )
-    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    -> (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
     (MLXArray, MLXArray, MLXArray, MLXArray)
 {
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless) {
@@ -530,3 +551,556 @@ public func compile(
         return (r[0], r[1], r[2], r[3])
     }
 }
+
+// MARK: - no state/Sendable
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray) -> (MLXArray, MLXArray)
+) -> @Sendable (MLXArray) -> (MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0])
+        return [r.0, r.1]
+    }
+    return { a in
+        let r = compileState.call([a])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray) -> (MLXArray, MLXArray, MLXArray)
+) -> @Sendable (MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0])
+        return [r.0, r.1, r.2]
+    }
+    return { a in
+        let r = compileState.call([a])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
+) -> @Sendable (MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a in
+        let r = compileState.call([a])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray)
+) -> @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1])
+        return [r.0, r.1]
+    }
+    return { a, b in
+        let r = compileState.call([a, b])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
+) -> @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b in
+        let r = compileState.call([a, b])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
+) -> @Sendable (MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b in
+        let r = compileState.call([a, b])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
+) -> @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2])
+        return [r.0, r.1]
+    }
+    return { a, b, c in
+        let r = compileState.call([a, b, c])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
+) -> @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b, c in
+        let r = compileState.call([a, b, c])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray, MLXArray
+        )
+) -> @Sendable (MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b, c in
+        let r = compileState.call([a, b, c])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        [f(args[0], args[1], args[2], args[3])]
+    }
+    return { a, b, c, d in
+        compileState.call([a, b, c, d])[0]
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3])
+        return [r.0, r.1]
+    }
+    return { a, b, c, d in
+        let r = compileState.call([a, b, c, d])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray
+        )
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b, c, d in
+        let r = compileState.call([a, b, c, d])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray, MLXArray
+        )
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray, MLXArray)
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b, c, d in
+        let r = compileState.call([a, b, c, d])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f: @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        [f(args[0], args[1], args[2], args[3], args[4])]
+    }
+    return { a, b, c, d, e in
+        compileState.call([a, b, c, d, e])[0]
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray
+        )
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray) {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4])
+        return [r.0, r.1]
+    }
+    return { a, b, c, d, e in
+        let r = compileState.call([a, b, c, d, e])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray
+        )
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray, MLXArray)
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b, c, d, e in
+        let r = compileState.call([a, b, c, d, e])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray, MLXArray
+        )
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        MLXArray, MLXArray, MLXArray, MLXArray
+    )
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b, c, d, e in
+        let r = compileState.call([a, b, c, d, e])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        [f(args[0], args[1], args[2], args[3], args[4], args[5])]
+    }
+    return { a, b, c, d, e, g in
+        compileState.call([a, b, c, d, e, g])[0]
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray
+        )
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (MLXArray, MLXArray)
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5])
+        return [r.0, r.1]
+    }
+    return { a, b, c, d, e, g in
+        let r = compileState.call([a, b, c, d, e, g])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray
+        )
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        MLXArray, MLXArray, MLXArray
+    )
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b, c, d, e, g in
+        let r = compileState.call([a, b, c, d, e, g])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+            MLXArray, MLXArray, MLXArray, MLXArray
+        )
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        MLXArray, MLXArray, MLXArray, MLXArray
+    )
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b, c, d, e, g in
+        let r = compileState.call([a, b, c, d, e, g])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> MLXArray
+) -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> MLXArray {
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])]
+    }
+    return { a, b, c, d, e, g, h in
+        compileState.call([a, b, c, d, e, g, h])[0]
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> (
+            MLXArray, MLXArray
+        )
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        MLXArray, MLXArray
+    )
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])
+        return [r.0, r.1]
+    }
+    return { a, b, c, d, e, g, h in
+        let r = compileState.call([a, b, c, d, e, g, h])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> (
+            MLXArray, MLXArray, MLXArray
+        )
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        MLXArray, MLXArray, MLXArray
+    )
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b, c, d, e, g, h in
+        let r = compileState.call([a, b, c, d, e, g, h])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray)
+        -> (
+            MLXArray, MLXArray, MLXArray, MLXArray
+        )
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) -> (
+        MLXArray, MLXArray, MLXArray, MLXArray
+    )
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b, c, d, e, g, h in
+        let r = compileState.call([a, b, c, d, e, g, h])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
+        -> MLXArray
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    MLXArray
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])]
+    }
+    return { a, b, c, d, e, g, h, i in
+        compileState.call([a, b, c, d, e, g, h, i])[0]
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
+        -> (MLXArray, MLXArray)
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    (MLXArray, MLXArray)
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])
+        return [r.0, r.1]
+    }
+    return { a, b, c, d, e, g, h, i in
+        let r = compileState.call([a, b, c, d, e, g, h, i])
+        return (r[0], r[1])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
+        -> (MLXArray, MLXArray, MLXArray)
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    (MLXArray, MLXArray, MLXArray)
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])
+        return [r.0, r.1, r.2]
+    }
+    return { a, b, c, d, e, g, h, i in
+        let r = compileState.call([a, b, c, d, e, g, h, i])
+        return (r[0], r[1], r[2])
+    }
+}
+
+@_documentation(visibility: internal)
+public func compile(
+    shapeless: Bool = false,
+    _ f:
+        @escaping @Sendable (
+            MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray
+        )
+        -> (MLXArray, MLXArray, MLXArray, MLXArray)
+)
+    -> @Sendable (MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray, MLXArray) ->
+    (MLXArray, MLXArray, MLXArray, MLXArray)
+{
+    let compileState = CompiledFunction(shapeless: shapeless) {
+        args in
+        let r = f(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])
+        return [r.0, r.1, r.2, r.3]
+    }
+    return { a, b, c, d, e, g, h, i in
+        let r = compileState.call([a, b, c, d, e, g, h, i])
+        return (r[0], r[1], r[2], r[3])
+    }
+}
diff --git a/Source/MLX/Transforms+Grad.swift b/Source/MLX/Transforms+Grad.swift
index e8d29e719..af1aecd8d 100644
--- a/Source/MLX/Transforms+Grad.swift
+++ b/Source/MLX/Transforms+Grad.swift
@@ -30,9 +30,7 @@ public func grad(
 // See ``grad(_:)-r8dv``
 public func grad(
     _ f: @escaping ([MLXArray]) -> MLXArray, argumentNumbers: some Collection<Int> = [0]
-) -> (
-    [MLXArray]
-) -> MLXArray {
+) -> ([MLXArray]) -> MLXArray {
     let wrappedFunction = wrapResult(wrapArguments(f))
     let gradientFunction = buildGradient(wrappedFunction, argumentNumbers: argumentNumbers)
     let uag: ([MLXArray]) -> [MLXArray] = unwrapArguments(gradientFunction)
diff --git a/Source/MLX/Transforms+Vmap.swift b/Source/MLX/Transforms+Vmap.swift
index d1cf92099..d2a463563 100644
--- a/Source/MLX/Transforms+Vmap.swift
+++ b/Source/MLX/Transforms+Vmap.swift
@@ -1,25 +1,27 @@
 import Cmlx
 import Foundation
 
-/// Returns a vectorized version of `f()`.
-///
-/// The returned function applies `f()` independently over the axis
-/// specified by `inAxes` and stacks the results along `outAxes`.
-///
-/// - Parameters:
-///   - f: Function to vectorize
-///   - inAxes: Axis of each input to map over. `nil` disables mapping for that input.
-///   - outAxes: Axis of each output to stack the results along
-/// - Returns: A vectorized function
-///
-/// ### See Also
-/// - <doc:vmap>
-public func vmap(
+final internal class UncheckedSendableBox<T>: @unchecked Sendable {
+    let value: T
+
+    init(_ value: T) {
+        self.value = value
+    }
+}
+
+/// Internal vmap implementation.  This takes a closure with sendability erased and produces
+/// a (declared) sendable closure.  Callers must declare their return sendability correctly.
+/// For example, ``vmap(_:inAxes:outAxes:)`` produces
+/// a non-Sendable closure from a non-Sendable input while
+/// ``vmapSendable(_:inAxes:outAxes:)`` produces
+/// Sendable from Sendable.
+private func vmapInternal(
     _ f: @escaping ([MLXArray]) -> [MLXArray],
-    inAxes: some Sequence<Int?> = [0],
-    outAxes: some Sequence<Int?> = [0]
-) -> ([MLXArray]) -> [MLXArray] {
-    { arrays in
+    inAxes: some Sequence<Int?> & Sendable = [0],
+    outAxes: some Sequence<Int?> & Sendable = [0]
+) -> @Sendable ([MLXArray]) -> [MLXArray] {
+    let box = UncheckedSendableBox(f)
+    return { arrays in
         let inAxes32 = inAxes.map { Int32($0 ?? -1) }
         let outAxes32 = outAxes.map { Int32($0 ?? -1) }
 
@@ -30,7 +32,7 @@ public func vmap(
         var traceOutputs = mlx_vector_array_new()
 
         withEvalLock {
-            let closure = new_mlx_closure(f)
+            let closure = new_mlx_closure(box.value)
             _ = inAxes32.withUnsafeBufferPointer { inAxesBuf in
                 mlx_detail_vmap_trace(
                     &traceInputs, &traceOutputs, closure, inputs, inAxesBuf.baseAddress,
@@ -66,6 +68,55 @@ public func vmap(
     }
 }
 
+/// Returns a vectorized version of `f()`.
+///
+/// The returned function applies `f()` independently over the axis
+/// specified by `inAxes` and stacks the results along `outAxes`.
+///
+/// - Parameters:
+///   - f: Function to vectorize
+///   - inAxes: Axis of each input to map over. `nil` disables mapping for that input.
+///   - outAxes: Axis of each output to stack the results along
+/// - Returns: A vectorized function
+///
+/// ### See Also
+/// - <doc:vmap>
+public func vmap(
+    _ f: @escaping ([MLXArray]) -> [MLXArray],
+    inAxes: some Sequence<Int?> & Sendable = [0],
+    outAxes: some Sequence<Int?> & Sendable = [0]
+) -> ([MLXArray]) -> [MLXArray] {
+    vmapInternal(f, inAxes: inAxes, outAxes: outAxes)
+}
+
+/// Returns a vectorized version of `f()`, a `Sendable` function (does not capture
+/// any non-Sendable state, such as an `MLXArray`).
+///
+/// The returned function applies `f()` independently over the axis
+/// specified by `inAxes` and stacks the results along `outAxes`.
+///
+/// ```swift
+/// @Sendable
+/// func add(_ x: MLXArray, _ y: MLXArray) -> MLXArray { x + y }
+/// let vf = vmapSendable(add, inAxes: (0, nil))
+/// ```
+///
+/// - Parameters:
+///   - f: Function to vectorize
+///   - inAxes: Axis of each input to map over. `nil` disables mapping for that input.
+///   - outAxes: Axis of each output to stack the results along
+/// - Returns: A vectorized function
+///
+/// ### See Also
+/// - <doc:vmap>
+public func vmapSendable(
+    _ f: @escaping @Sendable ([MLXArray]) -> [MLXArray],
+    inAxes: some Sequence<Int?> & Sendable = [0],
+    outAxes: some Sequence<Int?> & Sendable = [0]
+) -> @Sendable ([MLXArray]) -> [MLXArray] {
+    vmapInternal(f, inAxes: inAxes, outAxes: outAxes)
+}
+
 /// Overload of ``vmap(_:inAxes:outAxes:)`` for a single ``MLXArray`` input and
 /// output.
 ///
@@ -81,6 +132,20 @@ public func vmap(
     return { a in inner([a])[0] }
 }
 
+/// Overload of ``vmapSendable(_:inAxes:outAxes:)``
+/// for a single ``MLXArray`` input and output.
+///
+/// ### See Also
+/// - <doc:vmap>
+public func vmapSendable(
+    _ f: @escaping @Sendable (MLXArray) -> MLXArray,
+    inAxes: Int? = 0,
+    outAxes: Int? = 0
+) -> @Sendable (MLXArray) -> MLXArray {
+    let inner = vmapSendable({ [f($0[0])] }, inAxes: [inAxes], outAxes: [outAxes])
+    return { a in inner([a])[0] }
+}
+
 /// Overload of ``vmap(_:inAxes:outAxes:)`` for two ``MLXArray`` inputs and a
 /// single ``MLXArray`` output.
 ///
@@ -95,3 +160,18 @@ public func vmap(
     let inner = vmap({ [f($0[0], $0[1])] }, inAxes: [inAxes.0, inAxes.1], outAxes: [outAxes])
     return { a, b in inner([a, b])[0] }
 }
+
+/// Overload of ``vmapSendable(_:inAxes:outAxes:)``
+/// for a two ``MLXArray`` inputs and a single output.
+///
+/// ### See Also
+/// - <doc:vmap>
+public func vmapSendable(
+    _ f: @escaping @Sendable (MLXArray, MLXArray) -> MLXArray,
+    inAxes: (Int?, Int?) = (0, 0),
+    outAxes: Int? = 0
+) -> @Sendable (MLXArray, MLXArray) -> MLXArray {
+    let inner = vmapSendable(
+        { [f($0[0], $0[1])] }, inAxes: [inAxes.0, inAxes.1], outAxes: [outAxes])
+    return { a, b in inner([a, b])[0] }
+}
diff --git a/Tests/MLXTests/CompileLockOrderingTests.swift b/Tests/MLXTests/CompileLockOrderingTests.swift
index af7c67057..fd5378c1f 100644
--- a/Tests/MLXTests/CompileLockOrderingTests.swift
+++ b/Tests/MLXTests/CompileLockOrderingTests.swift
@@ -136,14 +136,14 @@ import XCTest
         /// compiler-cache erase -- happen at a chosen instant on a chosen
         /// thread.
         private final class Victim: @unchecked Sendable {
-            private var compiled: (@Sendable (MLXArray) -> MLXArray)?
+            private var compiled: ((MLXArray) -> MLXArray)?
 
             init() {
                 // A fresh compile() per instance, so the erase has a populated
                 // cache entry to remove.  The unique constant keeps this from
                 // sharing an entry with anything else.
                 let bias = MLXArray(Float.random(in: 1 ... 2))
-                compiled = MLX.compile { (x: MLXArray) in x * 3 + bias }
+                compiled = MLX.compile(inputs: [bias]) { (x: MLXArray) in x * 3 + bias }
             }
 
             /// Populate the cache entry that `deinit` will erase.
@@ -291,7 +291,8 @@ import XCTest
                 for i in 0 ..< rounds {
                     // A fresh compile each round keeps the cache cold, so every
                     // round really performs an insert.
-                    let compiled = MLX.compile { (x: MLXArray) in x * 2 + MLXArray(Float(i)) }
+                    let compiled = MLX.compile { (x: MLXArray) in x * 2 + MLXArray(Float(i))
+                    }
                     let r = compiled(MLXArray(Float(1)))
                     eval(r)
                     values.append(r.item(Float.self))
diff --git a/Tests/MLXTests/EvalTests.swift b/Tests/MLXTests/EvalTests.swift
index abaf825dd..234f67d83 100644
--- a/Tests/MLXTests/EvalTests.swift
+++ b/Tests/MLXTests/EvalTests.swift
@@ -613,32 +613,6 @@ struct EvalTests {
         #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.
diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift
index 5d1e3727a..6bc06ff66 100644
--- a/Tests/MLXTests/OptimizerTests.swift
+++ b/Tests/MLXTests/OptimizerTests.swift
@@ -67,6 +67,7 @@ class OptimizerTests: XCTestCase {
         // measure the distance from the prediction (model(x)) and the
         // ground truth (y).  this gives feedback on how close the
         // prediction is from matching the truth
+        @Sendable
         func loss(model: LinearFunctionModel, x: MLXArray, y: MLXArray) -> MLXArray {
             mseLoss(predictions: model(x), targets: y, reduction: .mean)
         }
diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift
index 71b58c990..629fd529e 100644
--- a/Tests/MLXTests/TransformTests.swift
+++ b/Tests/MLXTests/TransformTests.swift
@@ -112,6 +112,7 @@ class TransformTests: XCTestCase {
     }
 
     func testCompile() {
+        @Sendable
         func f(inputs: [MLXArray]) -> [MLXArray] {
             [square(inputs[0] * inputs[1])]
         }
@@ -231,11 +232,14 @@ class TransformTests: XCTestCase {
     }
 
     func testCompileCapturesModuleParameterWithoutState() {
-        // verify the capture path happens
+        // verify that compile captures all values used if not
+        // passed as inputs.
 
         let model = ScaleModule(2.0)
 
-        let compiled = compile { (x: MLXArray) -> MLXArray in
+        // lie about the inputs so that we can use a non-Sendable
+        // closure.  this is meant to demonstrate the failure
+        let compiled = compile(inputs: []) { (x: MLXArray) -> MLXArray in
             model(x)
         }
 
@@ -250,7 +254,7 @@ class TransformTests: XCTestCase {
         let r2 = compiled(MLXArray(Float(3)))
         XCTAssertEqual(
             r2.item(Float.self), 6,
-            "compiled function incorrectly ignored the updated module parameter")
+            "compiled function saw the update and it should not")
     }
 
     func testCompileWithStateObservesModuleParameterUpdates() {
@@ -276,6 +280,7 @@ class TransformTests: XCTestCase {
     }
 
     func testCompiledRandom() {
+        @Sendable
         func f(_ bias: MLXArray) -> MLXArray {
             MLXRandom.uniform(0 ..< 1, [4]) + bias
         }
@@ -335,7 +340,26 @@ class TransformTests: XCTestCase {
         assertEqual(vf(x), expected)
     }
 
+    func testVmapSendableSimple() async {
+        @Sendable
+        func f(_ x: MLXArray) -> MLXArray { x * 2 }
+
+        let x = MLXArray(0 ..< 6, [3, 2])
+
+        // this result type is Sendable
+        let vf = vmapSendable(f)
+
+        let t = Task {
+            let expected = stacked((0 ..< 3).map { f(x[$0]) }, axis: 0)
+
+            assertEqual(vf(x), expected)
+        }
+
+        _ = await t.value
+    }
+
     func testVmapTwoInputs() {
+        @Sendable
         func f(_ x: MLXArray, _ y: MLXArray) -> MLXArray { x + y }
 
         let x = MLXArray(0 ..< 6, [3, 2])
@@ -348,6 +372,7 @@ class TransformTests: XCTestCase {
     }
 
     func testVmapAxisPermutations() {
+        @Sendable
         func f(_ x: MLXArray) -> MLXArray { x * 2 }
 
         let x = MLXArray(0 ..< 6, [2, 3])
@@ -366,6 +391,7 @@ class TransformTests: XCTestCase {
     }
 
     func testVmapTwoInputsAxisPermutations() {
+        @Sendable
         func f(_ x: MLXArray, _ y: MLXArray) -> MLXArray { x + y }
 
         let x = MLXArray(0 ..< 6, [2, 3])
@@ -381,6 +407,7 @@ class TransformTests: XCTestCase {
     }
 
     func testVmapWithCompile() {
+        @Sendable
         func f(_ x: MLXArray) -> MLXArray { x.square() }
 
         let x = MLXRandom.normal([4, 2])
@@ -389,7 +416,9 @@ class TransformTests: XCTestCase {
         let mapCompile = vmap(compiled)
         assertEqual(mapCompile(x), vmap(f)(x))
 
-        let compiledMap = compile(vmap(f))
+        // use a non-Sendable version of compile because vmap doesn't produce
+        // a Sendable function
+        let compiledMap = compile(inputs: [], vmap(f))
         assertEqual(compiledMap(x), vmap(f)(x))
     }
 
@@ -408,6 +437,7 @@ class TransformTests: XCTestCase {
 
     func testCompileThreadSafety() async throws {
 
+        @Sendable
         func swiglu(_ xLinear: MLXArray, _ xGlu: MLXArray, alpha: Float = 1.702, limit: Float = 7.0)
             -> MLXArray
         {
diff --git a/Tests/MLXTests/WiredMemoryTests.swift b/Tests/MLXTests/WiredMemoryTests.swift
index 605e940b0..32a958fab 100644
--- a/Tests/MLXTests/WiredMemoryTests.swift
+++ b/Tests/MLXTests/WiredMemoryTests.swift
@@ -590,7 +590,7 @@ final class WiredMemoryTests: XCTestCase {
     /// Collects events from a stream until a predicate matches or a timeout fires.
     private static func collectEvents(
         stream: AsyncStream<WiredMemoryEvent>,
-        until predicate: @Sendable @escaping (WiredMemoryEvent) -> Bool,
+        until predicate: @escaping @Sendable (WiredMemoryEvent) -> Bool,
         timeout: TimeInterval = 10
     ) async throws -> [WiredMemoryEvent] {
         return try await withThrowingTaskGroup(of: [WiredMemoryEvent].self) { group in

From 17f64707b0aa2fca52e8192d4805e60e60fe1295 Mon Sep 17 00:00:00 2001
From: Purav Manot <60108184+pmanot@users.noreply.github.com>
Date: Mon, 19 Oct 2026 21:49:58 +0530
Subject: [PATCH 2/4] Fix iOS 26.2 launch crash from Logger.isEnabled (#494)

* Fix iOS 26.2 launch crash from Logger.isEnabled

* rebase and remove previous logging patch

---------

Co-authored-by: David Koski <dkoski@apple.com>
---
 .github/workflows/release.yml            |   3 +-
 CMakeLists.txt                           |  16 +-
 Package.swift                            |   7 -
 Source/MLX/Logging.swift                 |  14 +-
 Source/cLogSupport/cLogSupport.c         |   1 -
 Source/cLogSupport/include/cLogSupport.h |   9 -
 tools/create-xcframework.sh              |   2 +-
 xcode/MLX.xcodeproj/project.pbxproj      | 307 -----------------------
 8 files changed, 10 insertions(+), 349 deletions(-)
 delete mode 100644 Source/cLogSupport/cLogSupport.c
 delete mode 100644 Source/cLogSupport/include/cLogSupport.h

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 044a21a3e..127e29465 100644
--- a/.github/workflows/release.yml
+++ b/.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
diff --git a/CMakeLists.txt b/CMakeLists.txt
index f75ba6c9c..9c558a90c 100644
--- a/CMakeLists.txt
+++ b/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
diff --git a/Package.swift b/Package.swift
index 6188857d2..693f05a46 100644
--- a/Package.swift
+++ b/Package.swift
@@ -349,17 +349,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,
diff --git a/Source/MLX/Logging.swift b/Source/MLX/Logging.swift
index 174e09033..2c4245a73 100644
--- a/Source/MLX/Logging.swift
+++ b/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)")
             }
diff --git a/Source/cLogSupport/cLogSupport.c b/Source/cLogSupport/cLogSupport.c
deleted file mode 100644
index ae819f3c6..000000000
--- a/Source/cLogSupport/cLogSupport.c
+++ /dev/null
@@ -1 +0,0 @@
-void cLogSupport(void) { }
diff --git a/Source/cLogSupport/include/cLogSupport.h b/Source/cLogSupport/include/cLogSupport.h
deleted file mode 100644
index 9dc3f9262..000000000
--- a/Source/cLogSupport/include/cLogSupport.h
+++ /dev/null
@@ -1,9 +0,0 @@
-#if defined(__APPLE__)
-#import <os/log.h>
-#import <stdbool.h>
-
-static inline bool is_os_log_enabled(os_log_t log, os_log_type_t type) {
-    return os_log_type_enabled(log, type);
-}
-#endif
-
diff --git a/tools/create-xcframework.sh b/tools/create-xcframework.sh
index 676592543..f3fbb0726 100755
--- a/tools/create-xcframework.sh
+++ b/tools/create-xcframework.sh
@@ -20,7 +20,7 @@ YELLOW='\033[1;33m'
 NC='\033[0m' # No Color
 
 # Frameworks to package
-FRAMEWORKS=("MLX" "Cmlx" "cLogSupport" "MLXNN" "MLXOptimizers")
+FRAMEWORKS=("MLX" "Cmlx" "MLXNN" "MLXOptimizers")
 
 # Destinations to build for (ordered array)
 PLATFORM_NAMES=("iOS" "macOS" "Mac Catalyst" "tvOS" "visionOS")
diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj
index 00b532f54..710a87ff4 100644
--- a/xcode/MLX.xcodeproj/project.pbxproj
+++ b/xcode/MLX.xcodeproj/project.pbxproj
@@ -13,8 +13,6 @@
 		C35EA5643054AC8E00281871 /* MLXNN.framework in Embed Frameworks */ = {isa = PBXBuildFile; fileRef = C3CBE6B52EAC14DE0029A645 /* MLXNN.framework */; settings = {ATTRIBUTES = (CodeSignOnCopy, RemoveHeadersOnCopy, ); }; };
 		C35EA5673054AC8E00281871 /* MLXOptimizers.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C3CBE6CF2EAC15310029A645 /* MLXOptimizers.framework */; };
 		C35EA5683054AC8E00281871 /* MLXOptimizers.framework in Embed Frameworks */ = {isa = PBXBuildFile; fileRef = C3CBE6CF2EAC15310029A645 /* MLXOptimizers.framework */; settings = {ATTRIBUTES = (CodeSignOnCopy, RemoveHeadersOnCopy, ); }; };
-		C37BFFFA306C6138006C8EBC /* cLogSupport.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C37BFFEC306C60EE006C8EBC /* cLogSupport.framework */; platformFilters = (ios, maccatalyst, macos, xros, ); };
-		C37BFFFB306C6138006C8EBC /* cLogSupport.framework in Embed Frameworks */ = {isa = PBXBuildFile; fileRef = C37BFFEC306C60EE006C8EBC /* cLogSupport.framework */; settings = {ATTRIBUTES = (CodeSignOnCopy, RemoveHeadersOnCopy, ); }; };
 		C3AE8ECE2EAAA15F000BD280 /* MLX.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C3AE8EC42EAAA15F000BD280 /* MLX.framework */; };
 		C3AEA0322EAAE79A000BD280 /* Accelerate.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C3AEA0312EAAE79A000BD280 /* Accelerate.framework */; };
 		C3AEA0342EAAE7A0000BD280 /* Metal.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C3AEA0332EAAE7A0000BD280 /* Metal.framework */; };
@@ -49,13 +47,6 @@
 			remoteGlobalIDString = C3CBE6CE2EAC15310029A645;
 			remoteInfo = MLXOptimizers;
 		};
-		C37BFFFC306C6138006C8EBC /* PBXContainerItemProxy */ = {
-			isa = PBXContainerItemProxy;
-			containerPortal = C3AE8EBB2EAAA15F000BD280 /* Project object */;
-			proxyType = 1;
-			remoteGlobalIDString = C37BFFEB306C60EE006C8EBC;
-			remoteInfo = cLogSupport;
-		};
 		C3AE8ECF2EAAA15F000BD280 /* PBXContainerItemProxy */ = {
 			isa = PBXContainerItemProxy;
 			containerPortal = C3AE8EBB2EAAA15F000BD280 /* Project object */;
@@ -119,7 +110,6 @@
 			dstPath = "";
 			dstSubfolderSpec = 10;
 			files = (
-				C37BFFFB306C6138006C8EBC /* cLogSupport.framework in Embed Frameworks */,
 			);
 			name = "Embed Frameworks";
 			runOnlyForDeploymentPostprocessing = 0;
@@ -129,7 +119,6 @@
 /* Begin PBXFileReference section */
 		C319FC0A2F1564BE00620953 /* MLX.xctestplan */ = {isa = PBXFileReference; lastKnownFileType = text; path = MLX.xctestplan; sourceTree = "<group>"; };
 		C35E88D13054AACB00281871 /* MLXIntegrationTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = MLXIntegrationTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
-		C37BFFEC306C60EE006C8EBC /* cLogSupport.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = cLogSupport.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		C3AE8EC42EAAA15F000BD280 /* MLX.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = MLX.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		C3AE8ECD2EAAA15F000BD280 /* MLXTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = MLXTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		C3AE8EE62EAAA3C5000BD280 /* Cmlx.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Cmlx.framework; sourceTree = BUILT_PRODUCTS_DIR; };
@@ -169,17 +158,6 @@
 			);
 			target = C3AE8EC32EAAA15F000BD280 /* MLX */;
 		};
-		C37BFFF9306C6124006C8EBC /* Exceptions for "cLogSupport" folder in "cLogSupport" target */ = {
-			isa = PBXFileSystemSynchronizedBuildFileExceptionSet;
-			membershipExceptions = (
-				cLogSupport.c,
-				include/cLogSupport.h,
-			);
-			publicHeaders = (
-				include/cLogSupport.h,
-			);
-			target = C37BFFEB306C60EE006C8EBC /* cLogSupport */;
-		};
 		C3AE939F2EAAAA5A000BD280 /* Exceptions for "metal-cpp" folder in "Cmlx" target */ = {
 			isa = PBXFileSystemSynchronizedBuildFileExceptionSet;
 			membershipExceptions = (
@@ -1611,15 +1589,6 @@
 			path = ../Tests/MLXIntegrationTests;
 			sourceTree = SOURCE_ROOT;
 		};
-		C37BFFE4306C6075006C8EBC /* cLogSupport */ = {
-			isa = PBXFileSystemSynchronizedRootGroup;
-			exceptions = (
-				C37BFFF9306C6124006C8EBC /* Exceptions for "cLogSupport" folder in "cLogSupport" target */,
-			);
-			name = cLogSupport;
-			path = ../Source/cLogSupport;
-			sourceTree = SOURCE_ROOT;
-		};
 		C3AE8EDE2EAAA21C000BD280 /* xcconfig */ = {
 			isa = PBXFileSystemSynchronizedRootGroup;
 			path = xcconfig;
@@ -1731,19 +1700,11 @@
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
-		C37BFFE9306C60EE006C8EBC /* Frameworks */ = {
-			isa = PBXFrameworksBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-		};
 		C3AE8EC12EAAA15F000BD280 /* Frameworks */ = {
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
 				C3CBE67C2EABFBA60029A645 /* Numerics in Frameworks */,
-				C37BFFFA306C6138006C8EBC /* cLogSupport.framework in Frameworks */,
 				C3434FA12EABEF350010D6B8 /* Cmlx.framework in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
@@ -1791,7 +1752,6 @@
 		C3AE8EBA2EAAA15F000BD280 = {
 			isa = PBXGroup;
 			children = (
-				C37BFFE4306C6075006C8EBC /* cLogSupport */,
 				C319FC0A2F1564BE00620953 /* MLX.xctestplan */,
 				C3CBF17F2EAC22110029A645 /* ACKNOWLEDGMENTS.md */,
 				C3AE8EFC2EAAA40E000BD280 /* Cmlx */,
@@ -1821,7 +1781,6 @@
 				C3CBE6B52EAC14DE0029A645 /* MLXNN.framework */,
 				C3CBE6CF2EAC15310029A645 /* MLXOptimizers.framework */,
 				C35E88D13054AACB00281871 /* MLXIntegrationTests.xctest */,
-				C37BFFEC306C60EE006C8EBC /* cLogSupport.framework */,
 			);
 			name = Products;
 			sourceTree = "<group>";
@@ -1856,13 +1815,6 @@
 /* End PBXGroup section */
 
 /* Begin PBXHeadersBuildPhase section */
-		C37BFFE7306C60EE006C8EBC /* Headers */ = {
-			isa = PBXHeadersBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-		};
 		C3AE8EBF2EAAA15F000BD280 /* Headers */ = {
 			isa = PBXHeadersBuildPhase;
 			buildActionMask = 2147483647;
@@ -1920,26 +1872,6 @@
 			productReference = C35E88D13054AACB00281871 /* MLXIntegrationTests.xctest */;
 			productType = "com.apple.product-type.bundle.unit-test";
 		};
-		C37BFFEB306C60EE006C8EBC /* cLogSupport */ = {
-			isa = PBXNativeTarget;
-			buildConfigurationList = C37BFFF0306C60EE006C8EBC /* Build configuration list for PBXNativeTarget "cLogSupport" */;
-			buildPhases = (
-				C37BFFE7306C60EE006C8EBC /* Headers */,
-				C37BFFE8306C60EE006C8EBC /* Sources */,
-				C37BFFE9306C60EE006C8EBC /* Frameworks */,
-				C37BFFEA306C60EE006C8EBC /* Resources */,
-			);
-			buildRules = (
-			);
-			dependencies = (
-			);
-			name = cLogSupport;
-			packageProductDependencies = (
-			);
-			productName = cLogSupport;
-			productReference = C37BFFEC306C60EE006C8EBC /* cLogSupport.framework */;
-			productType = "com.apple.product-type.framework";
-		};
 		C3AE8EC32EAAA15F000BD280 /* MLX */ = {
 			isa = PBXNativeTarget;
 			buildConfigurationList = C3AE8ED62EAAA15F000BD280 /* Build configuration list for PBXNativeTarget "MLX" */;
@@ -1953,7 +1885,6 @@
 			buildRules = (
 			);
 			dependencies = (
-				C37BFFFD306C6138006C8EBC /* PBXTargetDependency */,
 			);
 			fileSystemSynchronizedGroups = (
 				C3434F7C2EABEF1E0010D6B8 /* MLX */,
@@ -2084,9 +2015,6 @@
 					C35E88D03054AACB00281871 = {
 						CreatedOnToolsVersion = 27.0;
 					};
-					C37BFFEB306C60EE006C8EBC = {
-						CreatedOnToolsVersion = 27.0;
-					};
 					C3AE8EC32EAAA15F000BD280 = {
 						CreatedOnToolsVersion = 16.4;
 					};
@@ -2127,7 +2055,6 @@
 				C3CBE6CE2EAC15310029A645 /* MLXOptimizers */,
 				C3AE8ECC2EAAA15F000BD280 /* MLXTests */,
 				C35E88D03054AACB00281871 /* MLXIntegrationTests */,
-				C37BFFEB306C60EE006C8EBC /* cLogSupport */,
 			);
 		};
 /* End PBXProject section */
@@ -2140,13 +2067,6 @@
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
-		C37BFFEA306C60EE006C8EBC /* Resources */ = {
-			isa = PBXResourcesBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-		};
 		C3AE8EC22EAAA15F000BD280 /* Resources */ = {
 			isa = PBXResourcesBuildPhase;
 			buildActionMask = 2147483647;
@@ -2192,13 +2112,6 @@
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
-		C37BFFE8306C60EE006C8EBC /* Sources */ = {
-			isa = PBXSourcesBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-		};
 		C3AE8EC02EAAA15F000BD280 /* Sources */ = {
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
@@ -2252,17 +2165,6 @@
 			target = C3CBE6CE2EAC15310029A645 /* MLXOptimizers */;
 			targetProxy = C35EA5693054AC8E00281871 /* PBXContainerItemProxy */;
 		};
-		C37BFFFD306C6138006C8EBC /* PBXTargetDependency */ = {
-			isa = PBXTargetDependency;
-			platformFilters = (
-				ios,
-				maccatalyst,
-				macos,
-				xros,
-			);
-			target = C37BFFEB306C60EE006C8EBC /* cLogSupport */;
-			targetProxy = C37BFFFC306C6138006C8EBC /* PBXContainerItemProxy */;
-		};
 		C3AE8ED02EAAA15F000BD280 /* PBXTargetDependency */ = {
 			isa = PBXTargetDependency;
 			target = C3AE8EC32EAAA15F000BD280 /* MLX */;
@@ -2448,206 +2350,6 @@
 			};
 			name = Release;
 		};
-		C37BFFF1306C60EE006C8EBC /* Debug */ = {
-			isa = XCBuildConfiguration;
-			buildSettings = {
-				ALLOW_TARGET_PLATFORM_SPECIALIZATION = YES;
-				ALWAYS_SEARCH_USER_PATHS = NO;
-				ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;
-				BUILD_LIBRARY_FOR_DISTRIBUTION = YES;
-				CLANG_ANALYZER_NONNULL = YES;
-				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
-				CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
-				CLANG_ENABLE_MODULES = YES;
-				CLANG_ENABLE_OBJC_ARC = YES;
-				CLANG_ENABLE_OBJC_WEAK = YES;
-				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
-				CLANG_WARN_BOOL_CONVERSION = YES;
-				CLANG_WARN_COMMA = YES;
-				CLANG_WARN_CONSTANT_CONVERSION = YES;
-				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
-				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
-				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
-				CLANG_WARN_EMPTY_BODY = YES;
-				CLANG_WARN_ENUM_CONVERSION = YES;
-				CLANG_WARN_INFINITE_RECURSION = YES;
-				CLANG_WARN_INT_CONVERSION = YES;
-				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
-				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
-				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
-				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
-				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
-				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
-				CLANG_WARN_STRICT_PROTOTYPES = YES;
-				CLANG_WARN_SUSPICIOUS_MOVE = YES;
-				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
-				CLANG_WARN_UNREACHABLE_CODE = YES;
-				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
-				CODE_SIGN_STYLE = Automatic;
-				COPY_PHASE_STRIP = NO;
-				CURRENT_PROJECT_VERSION = 1;
-				DEBUG_INFORMATION_FORMAT = dwarf;
-				DEFINES_MODULE = YES;
-				DYLIB_COMPATIBILITY_VERSION = 1;
-				DYLIB_CURRENT_VERSION = 1;
-				DYLIB_INSTALL_NAME_BASE = "@rpath";
-				ENABLE_MODULE_VERIFIER = YES;
-				ENABLE_STRICT_OBJC_MSGSEND = YES;
-				ENABLE_TESTABILITY = YES;
-				ENABLE_USER_SCRIPT_SANDBOXING = YES;
-				GCC_C_LANGUAGE_STANDARD = gnu17;
-				GCC_DYNAMIC_NO_PIC = NO;
-				GCC_NO_COMMON_BLOCKS = YES;
-				GCC_OPTIMIZATION_LEVEL = 0;
-				GCC_PREPROCESSOR_DEFINITIONS = (
-					"DEBUG=1",
-					"$(inherited)",
-				);
-				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
-				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
-				GCC_WARN_UNDECLARED_SELECTOR = YES;
-				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
-				GCC_WARN_UNUSED_FUNCTION = YES;
-				GCC_WARN_UNUSED_VARIABLE = YES;
-				GENERATE_INFOPLIST_FILE = YES;
-				INFOPLIST_KEY_NSHumanReadableCopyright = "Copyright © 2026 Apple Inc. All rights reserved.";
-				INSTALL_PATH = "$(LOCAL_LIBRARY_DIR)/Frameworks";
-				IPHONEOS_DEPLOYMENT_TARGET = 17.0;
-				LD_RUNPATH_SEARCH_PATHS = (
-					"@executable_path/Frameworks",
-					"@loader_path/Frameworks",
-				);
-				"LD_RUNPATH_SEARCH_PATHS[sdk=macosx*]" = (
-					"@executable_path/../Frameworks",
-					"@loader_path/Frameworks",
-				);
-				LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
-				MACOSX_DEPLOYMENT_TARGET = 14.0;
-				MARKETING_VERSION = 1.0;
-				MODULE_VERIFIER_SUPPORTED_LANGUAGES = "objective-c objective-c++";
-				MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20";
-				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
-				MTL_FAST_MATH = YES;
-				ONLY_ACTIVE_ARCH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = com.apple.cLogSupport;
-				PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)";
-				SDKROOT = auto;
-				SKIP_INSTALL = YES;
-				STRING_CATALOG_GENERATE_SYMBOLS = YES;
-				SUPPORTED_PLATFORMS = "appletvos appletvsimulator iphoneos iphonesimulator macosx xros xrsimulator";
-				SUPPORTS_MACCATALYST = YES;
-				SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)";
-				SWIFT_APPROACHABLE_CONCURRENCY = YES;
-				SWIFT_EMIT_LOC_STRINGS = YES;
-				SWIFT_INSTALL_MODULE = YES;
-				SWIFT_INSTALL_OBJC_HEADER = NO;
-				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
-				SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
-				SWIFT_VERSION = 5.0;
-				TARGETED_DEVICE_FAMILY = "1,2,3,7";
-				TVOS_DEPLOYMENT_TARGET = 17.0;
-				VERSIONING_SYSTEM = "apple-generic";
-				VERSION_INFO_PREFIX = "";
-				XROS_DEPLOYMENT_TARGET = 1.3;
-			};
-			name = Debug;
-		};
-		C37BFFF2306C60EE006C8EBC /* Release */ = {
-			isa = XCBuildConfiguration;
-			buildSettings = {
-				ALLOW_TARGET_PLATFORM_SPECIALIZATION = YES;
-				ALWAYS_SEARCH_USER_PATHS = NO;
-				ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;
-				BUILD_LIBRARY_FOR_DISTRIBUTION = YES;
-				CLANG_ANALYZER_NONNULL = YES;
-				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
-				CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
-				CLANG_ENABLE_MODULES = YES;
-				CLANG_ENABLE_OBJC_ARC = YES;
-				CLANG_ENABLE_OBJC_WEAK = YES;
-				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
-				CLANG_WARN_BOOL_CONVERSION = YES;
-				CLANG_WARN_COMMA = YES;
-				CLANG_WARN_CONSTANT_CONVERSION = YES;
-				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
-				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
-				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
-				CLANG_WARN_EMPTY_BODY = YES;
-				CLANG_WARN_ENUM_CONVERSION = YES;
-				CLANG_WARN_INFINITE_RECURSION = YES;
-				CLANG_WARN_INT_CONVERSION = YES;
-				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
-				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
-				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
-				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
-				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
-				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
-				CLANG_WARN_STRICT_PROTOTYPES = YES;
-				CLANG_WARN_SUSPICIOUS_MOVE = YES;
-				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
-				CLANG_WARN_UNREACHABLE_CODE = YES;
-				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
-				CODE_SIGN_STYLE = Automatic;
-				COPY_PHASE_STRIP = NO;
-				CURRENT_PROJECT_VERSION = 1;
-				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
-				DEFINES_MODULE = YES;
-				DYLIB_COMPATIBILITY_VERSION = 1;
-				DYLIB_CURRENT_VERSION = 1;
-				DYLIB_INSTALL_NAME_BASE = "@rpath";
-				ENABLE_MODULE_VERIFIER = YES;
-				ENABLE_NS_ASSERTIONS = NO;
-				ENABLE_STRICT_OBJC_MSGSEND = YES;
-				ENABLE_USER_SCRIPT_SANDBOXING = YES;
-				GCC_C_LANGUAGE_STANDARD = gnu17;
-				GCC_NO_COMMON_BLOCKS = YES;
-				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
-				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
-				GCC_WARN_UNDECLARED_SELECTOR = YES;
-				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
-				GCC_WARN_UNUSED_FUNCTION = YES;
-				GCC_WARN_UNUSED_VARIABLE = YES;
-				GENERATE_INFOPLIST_FILE = YES;
-				INFOPLIST_KEY_NSHumanReadableCopyright = "Copyright © 2026 Apple Inc. All rights reserved.";
-				INSTALL_PATH = "$(LOCAL_LIBRARY_DIR)/Frameworks";
-				IPHONEOS_DEPLOYMENT_TARGET = 17.0;
-				LD_RUNPATH_SEARCH_PATHS = (
-					"@executable_path/Frameworks",
-					"@loader_path/Frameworks",
-				);
-				"LD_RUNPATH_SEARCH_PATHS[sdk=macosx*]" = (
-					"@executable_path/../Frameworks",
-					"@loader_path/Frameworks",
-				);
-				LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
-				MACOSX_DEPLOYMENT_TARGET = 14.0;
-				MARKETING_VERSION = 1.0;
-				MODULE_VERIFIER_SUPPORTED_LANGUAGES = "objective-c objective-c++";
-				MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20";
-				MTL_ENABLE_DEBUG_INFO = NO;
-				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = com.apple.cLogSupport;
-				PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)";
-				SDKROOT = auto;
-				SKIP_INSTALL = YES;
-				STRING_CATALOG_GENERATE_SYMBOLS = YES;
-				SUPPORTED_PLATFORMS = "appletvos appletvsimulator iphoneos iphonesimulator macosx xros xrsimulator";
-				SUPPORTS_MACCATALYST = YES;
-				SWIFT_APPROACHABLE_CONCURRENCY = YES;
-				SWIFT_COMPILATION_MODE = wholemodule;
-				SWIFT_EMIT_LOC_STRINGS = YES;
-				SWIFT_INSTALL_MODULE = YES;
-				SWIFT_INSTALL_OBJC_HEADER = NO;
-				SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
-				SWIFT_VERSION = 5.0;
-				TARGETED_DEVICE_FAMILY = "1,2,3,7";
-				TVOS_DEPLOYMENT_TARGET = 17.0;
-				VERSIONING_SYSTEM = "apple-generic";
-				VERSION_INFO_PREFIX = "";
-				XROS_DEPLOYMENT_TARGET = 1.3;
-			};
-			name = Release;
-		};
 		C3AE8ED42EAAA15F000BD280 /* Debug */ = {
 			isa = XCBuildConfiguration;
 			baseConfigurationReferenceAnchor = C3AE8EDE2EAAA21C000BD280 /* xcconfig */;
@@ -2756,15 +2458,6 @@
 			defaultConfigurationIsVisible = 0;
 			defaultConfigurationName = Release;
 		};
-		C37BFFF0306C60EE006C8EBC /* Build configuration list for PBXNativeTarget "cLogSupport" */ = {
-			isa = XCConfigurationList;
-			buildConfigurations = (
-				C37BFFF1306C60EE006C8EBC /* Debug */,
-				C37BFFF2306C60EE006C8EBC /* Release */,
-			);
-			defaultConfigurationIsVisible = 0;
-			defaultConfigurationName = Release;
-		};
 		C3AE8EBE2EAAA15F000BD280 /* Build configuration list for PBXProject "MLX" */ = {
 			isa = XCConfigurationList;
 			buildConfigurations = (

From 12a704a6f6d10b682a91c0fbf7a6b01f20cccb9d Mon Sep 17 00:00:00 2001
From: Ronald Mannak <ronaldmannak@me.com>
Date: Mon, 19 Oct 2026 09:34:54 -0700
Subject: [PATCH 3/4] Add shardedToAllPredicate (#503)

---
 Source/MLXNN/Distributed.swift          | 34 +++++++++++++++++++++++--
 Tests/MLXTests/DistributedNNTests.swift | 34 +++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/Source/MLXNN/Distributed.swift b/Source/MLXNN/Distributed.swift
index 477c6088a..f401b09dd 100644
--- a/Source/MLXNN/Distributed.swift
+++ b/Source/MLXNN/Distributed.swift
@@ -217,6 +217,29 @@ private func quantizedShardedToAllPredicate(
     }
 }
 
+/// Shard the input dimensions of every parameter of a module that may hold
+/// quantized layers.
+///
+/// Integer positions count input features, so every ``QuantizedLinear`` gets
+/// them scaled to its own packed parameters, the way
+/// ``quantizedShardedToAllPredicate(_:inputDimensions:groupSize:size:)`` does
+/// for a single layer.  Every other parameter splits at the positions as given.
+private func shardedToAllPredicate(
+    _ segments: Segments, quantizedLayersIn module: Module, size: Int
+) throws -> ShardingPredicate {
+    var quantized = [String: ShardingPredicate]()
+    for (key, layer) in module.namedModules() {
+        guard let layer = layer as? QuantizedLinear else { continue }
+        quantized[key] = try quantizedShardedToAllPredicate(
+            segments, inputDimensions: layer.shape.1, groupSize: layer.groupSize, size: size)
+    }
+    let plain = shardedToAllPredicate(segments)
+    return { path, weight in
+        let owner = path.split(separator: ".").dropLast().joined(separator: ".")
+        return (quantized[owner] ?? plain)(path, weight)
+    }
+}
+
 /// Shard a module in place by replacing its parameters with sharded ones.
 ///
 /// The module itself is unchanged, so distributed communication only happens
@@ -231,8 +254,15 @@ public func shardInPlace(
     _ module: Module, sharding: ShardingType, segments: Segments = .count(1),
     group: MLXDistributed.Group? = nil
 ) throws {
-    try shardInPlace(
-        module, predicate: predicate(for: sharding, segments: segments), group: group)
+    let group = try group ?? MLXDistributed.initialize()
+    let sharded: ShardingPredicate
+    if case .shardedToAll = sharding, case .indices = segments {
+        sharded = try shardedToAllPredicate(
+            segments, quantizedLayersIn: module, size: group.size)
+    } else {
+        sharded = predicate(for: sharding, segments: segments)
+    }
+    try shardInPlace(module, predicate: sharded, group: group)
 }
 
 /// Shard a module in place, choosing how each parameter is split.
diff --git a/Tests/MLXTests/DistributedNNTests.swift b/Tests/MLXTests/DistributedNNTests.swift
index b57848678..861418bdf 100644
--- a/Tests/MLXTests/DistributedNNTests.swift
+++ b/Tests/MLXTests/DistributedNNTests.swift
@@ -584,6 +584,35 @@ private func assertGradientWhole(
         "\(key) gradient mismatch", file: file, line: line)
 }
 
+/// `shardInPlace` with integer segments on a module that holds a quantized
+/// layer.  The positions count input features, which the packed weight and the
+/// scales hold at different resolutions, so each has to split at its own index.
+func quantizedShardInPlaceBody(world: MLXDistributed.Group) throws {
+    MLXRandom.seed(0)
+    let x = MLXRandom.normal([4, 1024])
+    let model = Sequential(layers: QuantizedLinear(Linear(1024, 32, bias: false)))
+    let expected = model(x)
+
+    try shardInPlace(model, sharding: .shardedToAll, segments: .indices([512]), group: world)
+
+    // each process holds its slice of both segments; the sum of the partial
+    // results is the unsharded result
+    let step = 512 / world.size
+    let lower = world.rank * step
+    let xPart = concatenated(
+        [x[0..., lower ..< lower + step], x[0..., 512 + lower ..< 512 + lower + step]], axis: 1)
+    let result = MLXDistributed.allSum(model(xPart), group: world)
+    XCTAssertTrue(
+        expected.allClose(result, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+        "integer segments landed in the wrong place")
+
+    // a segment has to hold whole quantization groups on every process
+    XCTAssertThrowsError(
+        try shardInPlace(
+            QuantizedLinear(Linear(1024, 32)), sharding: .shardedToAll,
+            segments: .indices([528]), group: world))
+}
+
 /// The single process half: sharding is the identity in a group of size one,
 /// but the layers are still built and run.
 ///
@@ -625,6 +654,10 @@ class DistributedNNTests: XCTestCase {
     func testClipGradNormSharded() throws {
         try clipGradNormShardedBody(world: try MLXDistributed.initialize())
     }
+
+    func testQuantizedShardInPlace() throws {
+        try quantizedShardInPlaceBody(world: try MLXDistributed.initialize())
+    }
 }
 
 /// The multi process half, where every rank holds a different slice.
@@ -654,6 +687,7 @@ class DistributedNNRingTests: XCTestCase {
             try moduleSubstitutionBody(world: group)
             try donationBody(world: group)
             try quantizedShardedConstructionBody(world: group)
+            try quantizedShardInPlaceBody(world: group)
             try shardingEdgeCasesBody(world: group)
             try averageGradientsBody(world: group)
             try clipGradNormShardedBody(world: group)

From 863ee9f021ec9b82b4044f3aa7c03ec6db754a33 Mon Sep 17 00:00:00 2001
From: Shubhransh Gupta <54713516+shubhransh-gupta@users.noreply.github.com>
Date: Mon, 19 Oct 2026 22:49:26 +0530
Subject: [PATCH 4/4] Add Sendable closures and 3-parameter VJP support to
 MLXCustomFunction (#481)

* Add Sendable closures and 3-parameter VJP support to MLXCustomFunction (#475)
---
 Source/MLX/MLXCustomFunction.swift  | 52 +++++++++++++--------
 Tests/MLXTests/TransformTests.swift | 70 +++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 19 deletions(-)

diff --git a/Source/MLX/MLXCustomFunction.swift b/Source/MLX/MLXCustomFunction.swift
index 63eb66091..4d1c7f142 100644
--- a/Source/MLX/MLXCustomFunction.swift
+++ b/Source/MLX/MLXCustomFunction.swift
@@ -1,17 +1,30 @@
 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)
 }
 
@@ -19,8 +32,8 @@ 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()
 }
diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift
index 629fd529e..3e6b38c2e 100644
--- a/Tests/MLXTests/TransformTests.swift
+++ b/Tests/MLXTests/TransformTests.swift
@@ -561,4 +561,74 @@ class TransformTests: XCTestCase {
         }
     }
 
+    func testCustomFunctionVJP2() {
+        // Custom function with 2-parameter VJP: (primals, cotangents)
+        let customSquare = CustomFunction {
+            Forward { inputs in
+                [inputs[0].square()]
+            }
+            VJP { primals, cotangents in
+                // d/dx(x^2) = 2 * x, so cotangent * 2 * x
+                [cotangents[0] * 2 * primals[0]]
+            }
+        }
+
+        let x = MLXArray(3.0)
+        let forwardOut = customSquare([x])
+        assertEqual(forwardOut[0], MLXArray(9.0))
+
+        let gradFn = grad { inputs in
+            customSquare(inputs)
+        }
+        let grads = gradFn([x])
+        assertEqual(grads[0], MLXArray(6.0))
+    }
+
+    func testCustomFunctionVJP3WithOutputs() {
+        // Custom function with 3-parameter VJP: (primals, cotangents, outputs)
+        // For sigmoid: s(x) = output, and ds/dx = s * (1 - s) = output * (1 - output)
+        // This directly exercises the 3rd parameter `outputs`
+        let customSigmoid = CustomFunction {
+            Forward { inputs in
+                [sigmoid(inputs[0])]
+            }
+            VJP { primals, cotangents, outputs in
+                let out = outputs[0]
+                return [cotangents[0] * out * (1 - out)]
+            }
+        }
+
+        let x = MLXArray(0.0)
+        let forwardOut = customSigmoid([x])
+        assertEqual(forwardOut[0], MLXArray(0.5))
+
+        let gradFn = grad { inputs in
+            customSigmoid(inputs)
+        }
+        let grads = gradFn([x])
+        // At x = 0, sigmoid(0) = 0.5, grad = 0.5 * (1 - 0.5) = 0.25
+        assertEqual(grads[0], MLXArray(0.25))
+    }
+
+    func testCustomFunctionSendable() async {
+        // Verify Sendable closures and Sendable return from CustomFunction
+        let fn: @Sendable ([MLXArray]) -> [MLXArray] = CustomFunction {
+            Forward { inputs in
+                [inputs[0] + 1]
+            }
+            VJP { primals, cotangents, outputs in
+                [cotangents[0]]
+            }
+        }
+
+        // CustomFunction result is @Sendable and can be transferred across Task boundaries
+        let val = await Task { @Sendable in
+            let input = MLXArray(5.0)
+            let result = fn([input])
+            return result[0].item(Float.self)
+        }.value
+
+        XCTAssertEqual(val, 6.0)
+    }
+
 }