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/7] 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/7] 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/7] 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/7] 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)
+    }
+
 }

From d0e62b23f0bae51a7ca3d964f25cb8835e12a247 Mon Sep 17 00:00:00 2001
From: David Tai <dtaipublic@gmail.com>
Date: Tue, 20 Oct 2026 14:58:26 -0500
Subject: [PATCH 5/7] GPU.deviceInfo: select the device like MLX and compute it
 once (#498)

GPU.deviceInfo() called MTLCreateSystemDefaultDevice() on every call. In a
command-line process the first such call opens a window server session
(SLSCreateDefaultMetalDevice -> SLSMainDisplayID -> SLSInitialize ->
CGSDisplayInitialize). That step can stall for many seconds.

Select the device the same way as metal::load_device in MLX: the first
device of MTLCopyAllDevices(), else MTLCreateSystemDefaultDevice(). Compute
the value once per process. The public API and the returned fields do not
change.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
---
 Source/MLX/GPU+Metal.swift       | 20 ++++++++++++++++++--
 Tests/MLXTests/MemoryTests.swift | 15 +++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/Source/MLX/GPU+Metal.swift b/Source/MLX/GPU+Metal.swift
index cc4e83c89..975fbaa5f 100644
--- a/Source/MLX/GPU+Metal.swift
+++ b/Source/MLX/GPU+Metal.swift
@@ -242,13 +242,29 @@ public enum GPU {
     }
 
     /// Get information about the GPU device and system settings
+    ///
+    /// The value is computed once per process.
     public static func deviceInfo() -> DeviceInfo {
+        cachedDeviceInfo
+    }
+
+    private static let cachedDeviceInfo: DeviceInfo = {
         var mib = [CTL_HW, HW_MEMSIZE]
         var memSize: size_t = 0
         var length: size_t = MemoryLayout.size(ofValue: memSize)
         sysctl(&mib, 2, &memSize, &length, nil, 0)
 
-        if let device = MTLCreateSystemDefaultDevice() {
+        // Select the device the same way as MLX (metal::load_device): the first
+        // device of MTLCopyAllDevices(), else the system default device. In a
+        // command-line process MTLCreateSystemDefaultDevice() opens a window
+        // server session, which can stall for several seconds.
+        #if os(macOS)
+            let device = MTLCopyAllDevices().first ?? MTLCreateSystemDefaultDevice()
+        #else
+            let device = MTLCreateSystemDefaultDevice()
+        #endif
+
+        if let device {
             let architecture: String
             if #available(macOS 14.0, iOS 17.0, tvOS 17.0, *) {
                 architecture = device.architecture.name
@@ -265,5 +281,5 @@ public enum GPU {
                 architecture: "Unknown", maxBufferSize: 0, maxRecommendedWorkingSetSize: 0,
                 memorySize: memSize)
         }
-    }
+    }()
 }
diff --git a/Tests/MLXTests/MemoryTests.swift b/Tests/MLXTests/MemoryTests.swift
index 6e5120fa4..42485e76e 100644
--- a/Tests/MLXTests/MemoryTests.swift
+++ b/Tests/MLXTests/MemoryTests.swift
@@ -12,4 +12,19 @@ class MemoryTests: XCTestCase {
             print(x * x)
         }
     }
+
+    func testDeviceInfo() {
+        let info = GPU.deviceInfo()
+        XCTAssertGreaterThan(info.memorySize, 0)
+
+        guard Device.defaultDevice().deviceType == .gpu else { return }
+        XCTAssertNotEqual(info.architecture, "Unknown")
+        XCTAssertGreaterThan(info.maxBufferSize, 0)
+        XCTAssertGreaterThan(info.maxRecommendedWorkingSetSize, 0)
+
+        let again = GPU.deviceInfo()
+        XCTAssertEqual(again.architecture, info.architecture)
+        XCTAssertEqual(again.maxBufferSize, info.maxBufferSize)
+        XCTAssertEqual(again.maxRecommendedWorkingSetSize, info.maxRecommendedWorkingSetSize)
+    }
 }

From 9bd73fc77fe262b4b410d7e0bf2122eaf1099ae5 Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Tue, 20 Oct 2026 12:59:29 -0700
Subject: [PATCH 6/7] MaterializedArray is a Sendable MLXArray (#418)

* MaterializedArray is a Sendable MLXArray

- once an array has been evaluated we can use it as Sendable
- a subtype of MLXArray that is Sendable

- handle MaterializedArray passed as inputs (state) to a compiled function
	- contents are constant by definition so can be discarded and captured normally
	- calling materialize inside transforms
	- tracers inside materialized modules
---
 .../Articles/compilation.md                   |  30 +
 Source/MLX/Documentation.docc/MLX.md          |   9 +
 Source/MLX/Documentation.docc/MLXArray.md     |  12 +
 .../MLX/Documentation.docc/free-functions.md  |   5 +
 Source/MLX/MLXArray+Indexing.swift            |   3 +-
 Source/MLX/MLXArray+Init.swift                |  36 --
 Source/MLX/MLXArray.swift                     | 123 +++-
 Source/MLX/MaterializedArray.swift            | 137 ++++
 Source/MLX/Transforms+Compile.swift           |  76 ++-
 Source/MLX/Transforms+Eval.swift              |  81 +++
 Source/MLXNN/Documentation.docc/MLXNN.md      |  13 +
 Source/MLXNN/Documentation.docc/Module.md     |   1 +
 Source/MLXNN/Linear.swift                     |   4 +-
 Source/MLXNN/MaterializedModule.swift         | 202 ++++++
 Source/MLXNN/Module.swift                     | 185 +++++-
 Source/MLXNN/Quantized.swift                  |  31 +-
 .../MLXTests/LogicalParameterCountTests.swift | 132 ++++
 Tests/MLXTests/MLXArray+OpsTests.swift        |   2 +-
 Tests/MLXTests/MaterializedHazardTests.swift  | 597 ++++++++++++++++++
 Tests/MLXTests/MaterializedTests.swift        | 270 ++++++++
 skills/mlx-swift/SKILL.md                     |   5 +-
 skills/mlx-swift/references/concurrency.md    | 103 ++-
 22 files changed, 1957 insertions(+), 100 deletions(-)
 create mode 100644 Source/MLX/MaterializedArray.swift
 create mode 100644 Source/MLXNN/MaterializedModule.swift
 create mode 100644 Tests/MLXTests/LogicalParameterCountTests.swift
 create mode 100644 Tests/MLXTests/MaterializedHazardTests.swift
 create mode 100644 Tests/MLXTests/MaterializedTests.swift

diff --git a/Source/MLX/Documentation.docc/Articles/compilation.md b/Source/MLX/Documentation.docc/Articles/compilation.md
index b2f87e295..94cbc2044 100644
--- a/Source/MLX/Documentation.docc/Articles/compilation.md
+++ b/Source/MLX/Documentation.docc/Articles/compilation.md
@@ -345,6 +345,36 @@ This split matches how the result is used: a compiled function that captures
 mutable, non-`Sendable` state cannot safely be shared between threads, while a
 compiled pure function can.
 
+### Materialized state
+
+``MaterializedArray`` and `MaterializedModule` (from `MLXNN`) are `Sendable`, so
+a `@Sendable` closure may capture them and still use the `@Sendable` family:
+
+```swift
+let model = MaterializedModule(Linear(16, 16))
+
+// @Sendable -- the model's parameters are captured as constants
+let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in
+    model(x)
+}
+```
+
+Captured values are frozen at trace time, which is exactly right here: the
+contents of a ``MaterializedArray`` (and the parameters of a `MaterializedModule`)
+are sealed and can never change, so there is nothing to observe and no need for
+`inputs:`.
+
+What *can* change is which value the closure reaches.  If it captures a
+mutable container -- for example a `Mutex<MaterializedModule>` or a class that
+holds the current model -- and that container is later pointed at a different
+model, the compiled function keeps using the one it was traced with.  Compile a
+new function when the model is replaced.
+
+Sealed state reached through `inputs:` is also treated as a constant rather
+than being fed in on every call, but there the compiled function checks the
+identity of the sealed arrays on each call and retraces if they have been
+replaced.
+
 > Note: in the Swift 5 language mode with minimal concurrency checking the
 compiler does not diagnose non-`Sendable` captures, so a capturing closure
 literal may still resolve to the `@Sendable` variant. Use the Swift 6 language
diff --git a/Source/MLX/Documentation.docc/MLX.md b/Source/MLX/Documentation.docc/MLX.md
index fca5773a8..46b5ae120 100644
--- a/Source/MLX/Documentation.docc/MLX.md
+++ b/Source/MLX/Documentation.docc/MLX.md
@@ -92,6 +92,15 @@ See <doc:wired-memory> for configuration, best practices, and policy guidance.
 
 - ``MLXArray``
 
+### Concurrency
+
+`MLXArray` is not `Sendable`.  To share an array or model across task and actor
+boundaries, take an evaluated, immutable snapshot.
+
+- ``MaterializedArray``
+- ``materialize(_:)->MaterializedArray``
+- ``materialize(_:)->[MaterializedArray]``
+
 ### Free Functions
 
 - <doc:free-functions>
diff --git a/Source/MLX/Documentation.docc/MLXArray.md b/Source/MLX/Documentation.docc/MLXArray.md
index dd9a9875e..773901f68 100644
--- a/Source/MLX/Documentation.docc/MLXArray.md
+++ b/Source/MLX/Documentation.docc/MLXArray.md
@@ -42,6 +42,11 @@ let c = a + b
 may be large unresolved graphs) and the `addition` operator that combines them.  It is not safe
 to create `c` in one thread and consume/evaluate it in another.
 
+If you need to share an array across a task or actor boundary, take a
+``MaterializedArray`` snapshot via ``MLXArray/materialized()`` (or
+``materialize(_:)->MaterializedArray``).  A `MaterializedArray` is fully
+evaluated, immutable, and `Sendable`.
+
 ## Memory Safety
 
 > `MLXArray` is not memory safe.
@@ -91,3 +96,10 @@ and ``shape`` of the dimensions without changing the number of elements or conte
 ### Conversion and Data Types
 
 - <doc:conversion>
+
+### Concurrency
+
+- ``materialized()``
+- ``MaterializedArray``
+- ``materialize(_:)->MaterializedArray``
+- ``materialize(_:)->[MaterializedArray]``
diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md
index c54ebddbf..e0afc1fe5 100644
--- a/Source/MLX/Documentation.docc/free-functions.md
+++ b/Source/MLX/Documentation.docc/free-functions.md
@@ -219,6 +219,11 @@ operations as methods for convenience.
 - ``jvp(_:primals:tangents:)``
 - ``vjp(_:primals:cotangents:)``
 
+### Concurrency
+
+- ``materialize(_:)->MaterializedArray``
+- ``materialize(_:)->[MaterializedArray]``
+
 ### Device
 
 - ``using(device:fn:)``
diff --git a/Source/MLX/MLXArray+Indexing.swift b/Source/MLX/MLXArray+Indexing.swift
index bebaf23d2..394ab40d9 100644
--- a/Source/MLX/MLXArray+Indexing.swift
+++ b/Source/MLX/MLXArray+Indexing.swift
@@ -388,8 +388,7 @@ extension MLXArray {
                 var result = mlx_array_new()
                 mlx_scatter(
                     &result, self.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx)
-                mlx_array_set(&self.ctx, result)
-                mlx_array_free(result)
+                self._updateInternal(MLXArray(result))
                 return
             } else {
                 self._updateInternal(update)
diff --git a/Source/MLX/MLXArray+Init.swift b/Source/MLX/MLXArray+Init.swift
index 9aa04d109..b618a7950 100644
--- a/Source/MLX/MLXArray+Init.swift
+++ b/Source/MLX/MLXArray+Init.swift
@@ -683,39 +683,3 @@ extension MLXArray {
     }
 
 }
-
-// MARK: - Expressible by literals
-
-extension MLXArray: ExpressibleByArrayLiteral {
-
-    // Note: MLXArray does not implement ExpressibleByFloatLiteral etc. because
-    // we want to create arrays in the context of the other arrays.  For example:
-    //
-    // let x = MLXArray(1.5, dtype: .float16)
-    // let r = x + 2.5
-    //
-    // We expect r to have a dtype of float16.  See ``ScalarOrArray``.
-
-    /// Initializer allowing creation of 1d `MLXArray` from an array literal.
-    ///
-    /// ```swift
-    /// let a: MLXArray = [1, 2, 3]
-    /// ```
-    ///
-    /// This is convenient for methods that have `MLXArray` parameters:
-    ///
-    /// ```swift
-    /// print(array.take([1, 2, 3], axis: 0))
-    /// ```
-    ///
-    /// ### See Also
-    /// - <doc:initialization>
-    public convenience init(arrayLiteral elements: Int32...) {
-        let ctx = elements.withUnsafeBufferPointer { ptr in
-            let shape = [Int32(elements.count)]
-            return mlx_array_new_data(
-                ptr.baseAddress!, shape, Int32(shape.count), Int32.dtype.cmlxDtype)
-        }
-        self.init(ctx)
-    }
-}
diff --git a/Source/MLX/MLXArray.swift b/Source/MLX/MLXArray.swift
index bde1dd91a..1e87fe737 100644
--- a/Source/MLX/MLXArray.swift
+++ b/Source/MLX/MLXArray.swift
@@ -4,11 +4,25 @@ import Cmlx
 import Foundation
 import Numerics
 
-public final class MLXArray {
+public class MLXArray: ExpressibleByArrayLiteral {
 
     /// Internal pointer to the mlx-c wrapper on `mlx::core::array`, used with `Cmlx` interop.
     public internal(set) var ctx: mlx_array
 
+    /// True if the array is sealed (immutable): ``_updateInternal(_:)``, and so
+    /// every in-place mutation (`+=`, subscript assignment, the `compile`
+    /// state swap), will `fatalError`.
+    ///
+    /// Note: this is used for places where an MLXArray is materialized
+    /// but cannot be replaced with ``MaterializedArray``.  For example,
+    /// in a `Module`, `let x: MLXArray` can be updated via
+    /// `_updateInternal()` but can't be replaced with a
+    /// `MaterializedArray`.  This is used in that case.
+    ///
+    /// Code that has special handling for `MaterializedArrays`
+    /// (immutable) can use this property.
+    @_spi(Sealing) public internal(set) var isSealed = false
+
     /// Initialize with the given +1 context (transfer ownership).
     ///
     /// This initializer is for `Cmlx` interoperation.
@@ -19,6 +33,37 @@ public final class MLXArray {
         self.ctx = ctx
     }
 
+    // Note: MLXArray does not implement ExpressibleByFloatLiteral etc. because
+    // we want to create arrays in the context of the other arrays.  For example:
+    //
+    // let x = MLXArray(1.5, dtype: .float16)
+    // let r = x + 2.5
+    //
+    // We expect r to have a dtype of float16.  See ``ScalarOrArray``.
+
+    /// Initializer allowing creation of 1d `MLXArray` from an array literal.
+    ///
+    /// ```swift
+    /// let a: MLXArray = [1, 2, 3]
+    /// ```
+    ///
+    /// This is convenient for methods that have `MLXArray` parameters:
+    ///
+    /// ```swift
+    /// print(array.take([1, 2, 3], axis: 0))
+    /// ```
+    ///
+    /// ### See Also
+    /// - <doc:initialization>
+    required public convenience init(arrayLiteral elements: Int32...) {
+        let ctx = elements.withUnsafeBufferPointer { ptr in
+            let shape = [Int32(elements.count)]
+            return mlx_array_new_data(
+                ptr.baseAddress!, shape, Int32(shape.count), Int32.dtype.cmlxDtype)
+        }
+        self.init(ctx)
+    }
+
     /// return the equivalent of a `.none` MLXArray (for the C API).
     ///
     /// Not called `.none` to avoid ambiguity with `Optional`.  This can be used
@@ -37,7 +82,7 @@ public final class MLXArray {
     }
 
     /// Number of bytes per element
-    public var itemSize: Int { mlx_array_itemsize(ctx) }
+    final public var itemSize: Int { mlx_array_itemsize(ctx) }
 
     /// Total number of elements in the array
     ///
@@ -46,7 +91,7 @@ public final class MLXArray {
     /// print(array.size)
     /// // 12
     /// ```
-    public var size: Int { mlx_array_size(ctx) }
+    final public var size: Int { mlx_array_size(ctx) }
 
     /// Number of elements in the 0th dimension.
     ///
@@ -62,10 +107,10 @@ public final class MLXArray {
     ///     ...
     /// }
     /// ```
-    public var count: Int { dim(0) }
+    final public var count: Int { dim(0) }
 
     /// Number of bytes in the array.
-    public var nbytes: Int { mlx_array_nbytes(ctx) }
+    final public var nbytes: Int { mlx_array_nbytes(ctx) }
 
     /// Number of dimensions in the array.
     ///
@@ -74,7 +119,7 @@ public final class MLXArray {
     /// print(array.ndim)
     /// // 2
     /// ```
-    public var ndim: Int { mlx_array_ndim(ctx) }
+    final public var ndim: Int { mlx_array_ndim(ctx) }
 
     /// Data type of the elements in the array.
     ///
@@ -83,7 +128,7 @@ public final class MLXArray {
     /// print(array.dtype)
     /// // .int64 (aka Int.dtype)
     /// ```
-    public var dtype: DType { DType(mlx_array_dtype(ctx)) }
+    final public var dtype: DType { DType(mlx_array_dtype(ctx)) }
 
     /// Dimensions of the array.
     ///
@@ -92,7 +137,7 @@ public final class MLXArray {
     /// print(array.shape)
     /// // [3, 4]
     /// ```
-    public var shape: [Int] {
+    final public var shape: [Int] {
         let ndim = mlx_array_ndim(ctx)
         guard ndim > 0 else { return [] }
         let cShape = mlx_array_shape(ctx)!
@@ -104,7 +149,7 @@ public final class MLXArray {
     /// ```swift
     /// let (w, h) = array.shape2
     /// ```
-    public var shape2: (Int, Int) {
+    final public var shape2: (Int, Int) {
         let ndim = mlx_array_ndim(ctx)
         precondition(ndim == 2)
         let cShape = mlx_array_shape(ctx)!
@@ -116,7 +161,7 @@ public final class MLXArray {
     /// ```swift
     /// let (w, h, c) = array.shape3
     /// ```
-    public var shape3: (Int, Int, Int) {
+    final public var shape3: (Int, Int, Int) {
         let ndim = mlx_array_ndim(ctx)
         precondition(ndim == 3)
         let cShape = mlx_array_shape(ctx)!
@@ -128,7 +173,7 @@ public final class MLXArray {
     /// ```swift
     /// let (b, w, h, c) = array.shape4
     /// ```
-    public var shape4: (Int, Int, Int, Int) {
+    final public var shape4: (Int, Int, Int, Int) {
         let ndim = mlx_array_ndim(ctx)
         precondition(ndim == 4)
         let cShape = mlx_array_shape(ctx)!
@@ -149,7 +194,7 @@ public final class MLXArray {
     /// Strides of the array backing.
     ///
     /// Note: this is only stable once the array is evaluated.
-    var internalStrides: [Int] {
+    final var internalStrides: [Int] {
         let ndim = mlx_array_ndim(ctx)
         guard ndim > 0 else { return [] }
         let strides = mlx_array_strides(ctx)!
@@ -167,7 +212,7 @@ public final class MLXArray {
     /// // 4.5
     /// let value: Float = array[1].item()
     /// ```
-    public func item<T: HasDType>() -> T {
+    final public func item<T: HasDType>() -> T {
         item(T.self)
     }
 
@@ -328,7 +373,7 @@ public final class MLXArray {
     /// // 4.5
     /// let value = array[1].item(Float.self)
     /// ```
-    public func item<T: HasDType>(_ type: T.Type) -> T {
+    final public func item<T: HasDType>(_ type: T.Type) -> T {
         precondition(self.size == 1)
         eval()
 
@@ -466,7 +511,7 @@ public final class MLXArray {
     /// print(array.dim(1))
     /// // 4
     /// ```
-    public func dim(_ dim: Int) -> Int {
+    final public func dim(_ dim: Int) -> Int {
         Int(mlx_array_dim(ctx, MLX.resolve(axis: dim, ndim: mlx_array_ndim(ctx)).int32))
     }
 
@@ -481,7 +526,7 @@ public final class MLXArray {
     /// print(array.dim(index))
     /// // 4
     /// ```
-    func dim(_ dim: Int32) -> Int32 {
+    final func dim(_ dim: Int32) -> Int32 {
         mlx_array_dim(ctx, MLX.resolve(axis: Int(dim), ndim: mlx_array_ndim(ctx)).int32)
     }
 
@@ -492,7 +537,7 @@ public final class MLXArray {
     ///
     /// ### See Also
     /// - <doc:conversion>
-    public func asType(_ type: DType, stream: StreamOrDevice = .default) -> MLXArray {
+    final public func asType(_ type: DType, stream: StreamOrDevice = .default) -> MLXArray {
         guard type != self.dtype else { return self }
         var result = mlx_array_new()
         mlx_astype(&result, ctx, type.cmlxDtype, stream.ctx)
@@ -506,8 +551,9 @@ public final class MLXArray {
     ///
     /// ### See Also
     /// - <doc:conversion>
-    public func asType(_ type: (some HasDType).Type, stream: StreamOrDevice = .default) -> MLXArray
-    {
+    final public func asType(
+        _ type: (some HasDType).Type, stream: StreamOrDevice = .default
+    ) -> MLXArray {
         asType(type.dtype, stream: stream)
     }
 
@@ -524,7 +570,7 @@ public final class MLXArray {
     /// - ``realPart(stream:)``
     /// - ``imaginaryPart(stream:)``
     /// - <doc:conversion>
-    public func asImaginary(stream: StreamOrDevice = .default) -> MLXArray {
+    final public func asImaginary(stream: StreamOrDevice = .default) -> MLXArray {
         precondition(!dtype.isComplex)
         let i = MLXArray(real: 0, imaginary: 1)
         return self * i
@@ -534,7 +580,7 @@ public final class MLXArray {
     ///
     /// ### See Also
     /// - <doc:conversion>
-    public func realPart(stream: StreamOrDevice = .default) -> MLXArray {
+    final public func realPart(stream: StreamOrDevice = .default) -> MLXArray {
         precondition(dtype.isComplex)
         return asType(Float.self)
     }
@@ -543,7 +589,7 @@ public final class MLXArray {
     ///
     /// ### See Also
     /// - <doc:conversion>
-    public func imaginaryPart(stream: StreamOrDevice = .default) -> MLXArray {
+    final public func imaginaryPart(stream: StreamOrDevice = .default) -> MLXArray {
         precondition(dtype.isComplex)
         let i = MLXArray(real: 0, imaginary: 1)
         return (self / i).asType(.float32)
@@ -557,11 +603,36 @@ public final class MLXArray {
         MLX.eval(self)
     }
 
+    /// Force evaluation and return a ``MaterializedArray`` snapshot of `self`.
+    ///
+    /// The returned array is fully evaluated, immutable, and `Sendable` — see
+    /// ``MaterializedArray`` for the guarantees and trade-offs.  Use this when
+    /// you need to hand an array across task or actor boundaries, or when you
+    /// need a stable reference that cannot be mutated by other code holding
+    /// the original `MLXArray`.
+    ///
+    /// This is a thin convenience over ``materialize(_:)->MaterializedArray``.
+    ///
+    /// - Important: this must not be called inside a function transformation
+    ///   (`grad`, `vjp`, `compile`, etc.) and will `fatalError` if it is.
+    ///   See ``materialize(_:)->[MaterializedArray]``.
+    ///
+    /// ### See Also
+    /// - ``MaterializedArray``
+    /// - ``materialize(_:)->MaterializedArray``
+    public func materialized() -> MaterializedArray {
+        MLX.materialize(self)
+    }
+
     /// Replace the contents with a reference to a new array (INTERNAL).
     ///
     /// Note: this is an implementation detail and only visible because of the need to call it from
     /// other `mlx-swift` modules.
     public func _updateInternal(_ array: MLXArray) {
+        if isSealed {
+            fatalError("MLXArray is sealed (immutable) and cannot be updated")
+        }
+
         // Free the old array rather than assigning over it: `mlx_array_set` is a C++ copy assignment,
         // which skips `array::~array()` and with it the breaking of sibling cycles. An unevaluated
         // output of a multi-output primitive (e.g. the `Quantize` behind a freshly built
@@ -573,6 +644,14 @@ public final class MLXArray {
         self.ctx = new
     }
 
+    /// Seal the array: see ``isSealed``.  This is irreversible and idempotent.
+    ///
+    /// The array should already be evaluated, e.g. the parameters of a module
+    /// after `eval`.
+    @_spi(Sealing) public func _seal() {
+        isSealed = true
+    }
+
     /// Internal function for copying the backing `mlx::core::array` context.
     func copyContext() -> MLXArray {
         var new = mlx_array_new()
diff --git a/Source/MLX/MaterializedArray.swift b/Source/MLX/MaterializedArray.swift
new file mode 100644
index 000000000..3ed115937
--- /dev/null
+++ b/Source/MLX/MaterializedArray.swift
@@ -0,0 +1,137 @@
+// Copyright © 2026 Apple Inc.
+
+import Cmlx
+import Foundation
+import Numerics
+
+/// A fully-evaluated, immutable ``MLXArray`` that is safe to share across
+/// concurrency domains.
+///
+/// `MLXArray` is normally lazy: it is a handle to a node in a computation
+/// graph that is not realized until something forces it (a scalar read,
+/// an ``eval(_:)-(MLXArray...)`` call, etc.).  These unrealized arrays are
+/// not thread safe -- they require mutation for their evaluation.
+///
+/// `MaterializedArray` is a snapshot that closes that gap:
+///
+/// - The contents are evaluated at the moment of construction, so no further
+///   graph work is pending.
+/// - You can only create an instance via ``MLXArray/materialized()``
+///   or ``materialize(_:)->MaterializedArray``, and not inside a function
+///   transformation such as `grad` or `compile` (that traps, see
+///   ``materialize(_:)->[MaterializedArray]``)
+/// - Mutation methods are marked as unavailable and will `fatalError`
+///   if you somehow manage to call them (e.g. through dynamic types)
+/// - It is declared `@unchecked Sendable` and may be passed freely between
+///   tasks, actors, and other concurrency boundaries.
+///
+/// Construction is intentionally narrow.  Obtain one via:
+///
+/// ```swift
+/// let m1 = a.materialized()
+/// let m2 = materialize(a)
+/// ```
+///
+/// A `MaterializedArray` is itself an ``MLXArray`` and can be used anywhere
+/// an `MLXArray` is accepted.  Operations involving one still produce
+/// ordinary (lazy) `MLXArray` results — only the snapshot itself is frozen.
+///
+/// ### Mutation through an `MLXArray` reference
+///
+/// The in-place operators are unavailable when the static type is
+/// `MaterializedArray`, so this is a compile error:
+///
+/// ```swift
+/// var m = a.materialized()
+/// m += 1      // error: unavailable
+/// ```
+///
+/// However, `MLXArray`'s compound assignment operators (`+=`, etc.) and
+/// subscript assignment mutate the array in place.  If a `MaterializedArray`
+/// is held in a variable typed as `MLXArray` these still compile, but they
+/// trap at runtime rather than mutate the shared snapshot:
+///
+/// ```swift
+/// var w: MLXArray = a.materialized()
+/// w += 1      // compiles, traps at runtime
+/// w[0] = 0    // compiles, traps at runtime
+/// ```
+///
+/// The same applies to an operation that returns its input unchanged, e.g.
+/// `asType()` to the array's existing dtype.  To modify the values, produce
+/// a new array rather than mutating in place, e.g. `w = w + 1` rather
+/// than `w += 1`.
+///
+/// ### See Also
+/// - ``MLXArray/materialized()``
+/// - ``materialize(_:)->MaterializedArray``
+public final class MaterializedArray: MLXArray, @unchecked Sendable {
+
+    init(materialized ctx: consuming mlx_array) {
+        super.init(ctx)
+        self.isSealed = true
+    }
+
+    @available(*, unavailable, message: "MaterializedArray can only be created via materialize()")
+    required public convenience init(arrayLiteral elements: Int32...) {
+        fatalError("unavailable")
+    }
+
+    final public override func materialized() -> MaterializedArray {
+        self
+    }
+
+    // MARK: - Update sealing
+
+    @available(*, unavailable)
+    override public func _updateInternal(_ array: MLXArray) {
+        // Note that this might be called via:
+        //
+        // a[1] = b
+        // a += b
+        fatalError("unavailable")
+    }
+
+    // MARK: - In place operators
+
+    @available(*, unavailable)
+    public static func += (lhs: inout MaterializedArray, rhs: MLXArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func += (lhs: inout MaterializedArray, rhs: some ScalarOrArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func -= (lhs: inout MaterializedArray, rhs: MLXArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func -= (lhs: inout MaterializedArray, rhs: some ScalarOrArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func *= (lhs: inout MaterializedArray, rhs: MLXArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func *= (lhs: inout MaterializedArray, rhs: some ScalarOrArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func /= (lhs: inout MaterializedArray, rhs: MLXArray) {
+        fatalError("unavailable")
+    }
+
+    @available(*, unavailable)
+    public static func /= (lhs: inout MaterializedArray, rhs: some ScalarOrArray) {
+        fatalError("unavailable")
+    }
+
+}
diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift
index 527a82e44..d04944dfe 100644
--- a/Source/MLX/Transforms+Compile.swift
+++ b/Source/MLX/Transforms+Compile.swift
@@ -3,7 +3,9 @@
 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 --
+// except for ``CompiledFunction/tracedConstants``, which is guarded by
+// ``CompiledFunction/lock``.
 // 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:)
@@ -30,6 +32,27 @@ final class CompiledFunction: @unchecked (Sendable) {
 
     let shapeless: Bool
 
+    /// The sealed arrays reachable from ``inputs`` when the cached graphs were
+    /// traced, or `nil` before the first call.
+    ///
+    /// Sealed arrays are not threaded through the compiled function as inputs
+    /// (that would mean swapping tracers into them), they are baked into the
+    /// traced graph as constants.  The compiler cache only keys on the inputs,
+    /// so if the state is later rebound to *different* sealed arrays (a
+    /// different `MaterializedModule`, a replaced child module, a new
+    /// `MaterializedArray`) the cache would happily return a graph holding the
+    /// old ones.  ``innerCall(_:)`` compares the current sealed arrays against
+    /// these and erases the cache (forcing a retrace) if they differ.
+    ///
+    /// Identity (`===`) is sufficient: a sealed array's contents can never be
+    /// replaced, so the same object always refers to the same node.  The arrays
+    /// are retained so that a freed one can't be replaced by a new array at the
+    /// same address.  This costs no extra memory: the cached graph holds the
+    /// same data.
+    ///
+    /// Guarded by ``lock``.
+    private var tracedConstants: [MLXArray]?
+
     init(
         inputs: [any Updatable], outputs: [any Updatable], shapeless: Bool,
         _ f: @escaping ([MLXArray]) -> [MLXArray]
@@ -111,14 +134,20 @@ final class CompiledFunction: @unchecked (Sendable) {
                         + "see CompiledFunction.deinit")
                 EvalLockOwnership.recordCompileErase()
             #endif
-            // remove the compiled structure from the back end
-            var cache = mlx_compile_cache_new()
-            mlx_detail_compile_cache(&cache)
-            defer { mlx_compile_cache_free(cache) }
-            mlx_detail_compile_erase(cache, functionID)
+            Self.eraseCache(functionID)
         }
     }
 
+    /// Remove the compiled structure for `functionID` from the back end.
+    ///
+    /// - Precondition: `evalLock` is held -- see ``deinit``.
+    private static func eraseCache(_ functionID: UInt) {
+        var cache = mlx_compile_cache_new()
+        mlx_detail_compile_cache(&cache)
+        defer { mlx_compile_cache_free(cache) }
+        mlx_detail_compile_erase(cache, functionID)
+    }
+
     /// Evaluate the compiled function.
     ///
     /// ### Lock ordering
@@ -156,7 +185,8 @@ final class CompiledFunction: @unchecked (Sendable) {
             // ``lock`` exists only to guard the observed state (#226): the
             // `_updateInternal` swaps in `innerCall` are the sole mutation of
             // anything reachable from `self` -- every other stored property is
-            // `let`, and `id` is written once at init.  With no state there is
+            // `let` except ``tracedConstants`` (only written when there is
+            // state), and `id` is written once at init.  With no state there is
             // nothing to guard, and `evalLock` above already serializes the
             // compiler cache and every other caller.
             //
@@ -190,9 +220,33 @@ final class CompiledFunction: @unchecked (Sendable) {
     ///   too unless this function has no observed state; see ``call(_:)``,
     ///   which is the only caller.
     func innerCall(_ arguments: [MLXArray]) -> [MLXArray] {
-        let stateInputs = inputs.flatMap { $0.innerState() }
+        // Sealed arrays (MaterializedArray, the parameters of a sealed module)
+        // can never change, and must not be mutated, so they are not threaded
+        // through the tracer swap below -- they are simply captured as
+        // constants, the same as any other array reachable from the closure
+        // that isn't part of `inputs`/`outputs`.
+        let state = inputs.flatMap { $0.innerState() }
+        let stateInputs = state.filter { !$0.isSealed }
         let argumentsCount = arguments.count
 
+        // The sealed arrays are baked into the cached graphs, so if they have
+        // been rebound since those were traced the cache is stale -- see
+        // ``tracedConstants``.
+        if !inputs.isEmpty {
+            let constants = state.filter { $0.isSealed }
+            if let traced = tracedConstants {
+                let same =
+                    traced.count == constants.count
+                    && zip(traced, constants).allSatisfy { $0 === $1 }
+                if !same {
+                    Self.eraseCache(id)
+                    tracedConstants = constants
+                }
+            } else {
+                tracedConstants = constants
+            }
+        }
+
         // inner function to hande the compilation.  this is called
         // once per compile (typically once overall, but can be called
         // again if the conditions for recompile change)
@@ -217,7 +271,9 @@ final class CompiledFunction: @unchecked (Sendable) {
             let result = f(tracerArguments)
 
             // recapture the state as it may have changed
-            let stateOutputTracers = outputs.flatMap { $0.innerState() }.map { $0.copyContext() }
+            let stateOutputTracers = outputs.flatMap { $0.innerState() }
+                .filter { !$0.isSealed }
+                .map { $0.copyContext() }
 
             // put the original values back in the state
             for (s, saved) in zip(stateInputs, savedStateInputs) {
@@ -267,7 +323,7 @@ final class CompiledFunction: @unchecked (Sendable) {
         let resultsPlusStateOutput = mlx_vector_array_values(resultVector)
 
         // push the stateOutput into the state
-        let stateOutput = outputs.flatMap { $0.innerState() }
+        let stateOutput = outputs.flatMap { $0.innerState() }.filter { !$0.isSealed }
 
         for (s, newValues) in zip(stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) {
             s._updateInternal(newValues)
diff --git a/Source/MLX/Transforms+Eval.swift b/Source/MLX/Transforms+Eval.swift
index ac7b430f0..efc8c4d21 100644
--- a/Source/MLX/Transforms+Eval.swift
+++ b/Source/MLX/Transforms+Eval.swift
@@ -157,6 +157,87 @@ func isEvaluated(_ array: MLXArray) -> Bool {
     return available
 }
 
+/// Evaluate `array` and return a ``MaterializedArray`` snapshot of its contents.
+///
+/// The returned array is fully evaluated, immutable, and `Sendable`, so it can
+/// be passed across task and actor boundaries.  See ``MaterializedArray`` for
+/// the full set of guarantees.
+///
+/// `array` itself is unaffected — it remains the same lazy ``MLXArray`` it was
+/// before the call — but because evaluation is forced, any pending graph work
+/// behind it has been realized as a side effect.
+///
+/// ### Inside function transformations
+///
+/// `materialize` will `fatalError` if called inside a function transformation
+/// (`grad`, `valueAndGrad`, `vjp`, `jvp`, `vmap`, `compile`), including from
+/// a custom function (``CustomFunction(_:)``) body or VJP that is being run by
+/// one.  Inside `grad`/`vjp` an `eval` of a traced value computes it but
+/// does not detach it from the graph (the graph is needed for the backward
+/// pass), so the result could not honor the ``MaterializedArray`` contract:
+/// it would still reference the trace, and that graph is detached -- mutated
+/// -- by a later `eval` once tracing ends.  Inside `compile`/`vmap` traced
+/// values cannot be evaluated at all.
+///
+/// This is detected by the current thread holding the eval lock, which is
+/// the case for the whole duration of a transformation, so it also rejects
+/// arrays that do not depend on the traced inputs.  Materialize those before
+/// entering the transformation.
+///
+/// ### See Also
+/// - ``MaterializedArray``
+/// - ``MLXArray/materialized()``
+/// - ``materialize(_:)``
+public func materialize(_ array: MLXArray) -> MaterializedArray {
+    requireOutsideTransform()
+    eval(array)
+    var m = mlx_array_new()
+    mlx_array_set(&m, array.ctx)
+    return MaterializedArray(materialized: m)
+}
+
+/// Evaluate `arrays` and return a ``MaterializedArray`` snapshot for each one.
+///
+/// Equivalent to calling ``materialize(_:)`` on each element, but evaluates the
+/// whole batch in a single ``eval(_:)-(Sequence<Any>)`` call so MLX can schedule the
+/// work together.  Prefer this overload when you need to materialize several
+/// arrays at once — for example, the parameters of a model.
+///
+/// - Important: this must not be called inside a function transformation
+///   (`grad`, `valueAndGrad`, `vjp`, `jvp`, `vmap`, `compile`) and will
+///   `fatalError` if it is.  See
+///   ``materialize(_:)->MaterializedArray``.
+///
+/// ### See Also
+/// - ``MaterializedArray``
+/// - ``materialize(_:)``
+public func materialize(_ arrays: [MLXArray]) -> [MaterializedArray] {
+    requireOutsideTransform()
+    eval(arrays)
+    return arrays.map {
+        var m = mlx_array_new()
+        mlx_array_set(&m, $0.ctx)
+        return MaterializedArray(materialized: m)
+    }
+}
+
+/// Trap if called inside a function transformation -- see
+/// ``materialize(_:)->[MaterializedArray]``.
+///
+/// This is the `.nested` case of ``withEvalLockState(_:)``.  It reads the
+/// (thread-local) ownership directly rather than calling `withEvalLockState`:
+/// acquiring the lock just to check would add contention, and wrapping the
+/// `eval` in it would make that `eval` nested too, losing the wait outside of
+/// the lock in ``eval(_:)-(Collection<MLXArray>)``.
+private func requireOutsideTransform(_ caller: StaticString = #function) {
+    if EvalLockOwnership.isHeldByCurrentThread {
+        fatalError(
+            "\(caller) called inside a function transformation (grad, vjp, compile, etc.). "
+                + "A traced value cannot be materialized -- materialize before "
+                + "entering the transformation.")
+    }
+}
+
 /// Evaluate one or more `MLXArray`
 ///
 /// ### See Also
diff --git a/Source/MLXNN/Documentation.docc/MLXNN.md b/Source/MLXNN/Documentation.docc/MLXNN.md
index c527ebaf2..619fc5cb7 100644
--- a/Source/MLXNN/Documentation.docc/MLXNN.md
+++ b/Source/MLXNN/Documentation.docc/MLXNN.md
@@ -41,6 +41,18 @@ respect to these trainable parameters.
 
 See <doc:training>
 
+### Concurrency
+
+A ``Module`` is not `Sendable` because its parameters are lazy, mutable
+`MLXArray` values.  To share a model across task or actor boundaries, wrap it
+in a ``MaterializedModule``, which evaluates every parameter, replaces the
+mutable ones with `MaterializedArray` snapshots, and seals the module against
+further mutation.
+
+```swift
+let model = MaterializedModule(Linear(10, 10))
+```
+
 ## Other MLX Packages
 
 - [MLX](mlx)
@@ -58,6 +70,7 @@ See <doc:training>
 ### Base Classes and Interfaces
 
 - ``Module``
+- ``MaterializedModule``
 - ``UnaryLayer``
 - ``Quantizable``
 
diff --git a/Source/MLXNN/Documentation.docc/Module.md b/Source/MLXNN/Documentation.docc/Module.md
index ae940c4e2..54401e054 100644
--- a/Source/MLXNN/Documentation.docc/Module.md
+++ b/Source/MLXNN/Documentation.docc/Module.md
@@ -6,6 +6,7 @@
 
 - ``Module/apply(filter:map:)``
 - ``Module/filterMap(filter:map:isLeaf:)``
+- ``Module/logicalParameterCount``
 - ``Module/mapParameters(map:isLeaf:)``
 - ``Module/parameters()``
 - ``Module/trainableParameters()``
diff --git a/Source/MLXNN/Linear.swift b/Source/MLXNN/Linear.swift
index dd52c4226..df2467660 100644
--- a/Source/MLXNN/Linear.swift
+++ b/Source/MLXNN/Linear.swift
@@ -69,8 +69,8 @@ public class Identity: Module, UnaryLayer {
 /// - ``Bilinear``
 open class Linear: Module, UnaryLayer, Quantizable {
 
-    public let weight: MLXArray
-    public let bias: MLXArray?
+    @ParameterInfo public var weight: MLXArray
+    @ParameterInfo public var bias: MLXArray?
 
     open var shape: (Int, Int) {
         weight.shape2
diff --git a/Source/MLXNN/MaterializedModule.swift b/Source/MLXNN/MaterializedModule.swift
new file mode 100644
index 000000000..c2867f12f
--- /dev/null
+++ b/Source/MLXNN/MaterializedModule.swift
@@ -0,0 +1,202 @@
+// Copyright © 2026 Apple Inc.
+
+import MLX
+
+/// A container for a `Module` whose parameters have been materialized so that
+/// the whole module is safe to share across concurrency domains.
+///
+/// A normal ``Module`` is not `Sendable`: its parameters are `MLXArray`
+/// instances, which are lazy and may be mutated in place during evaluation
+/// or training.  `MaterializedModule` wraps a base module, evaluates every
+/// parameter, replaces each with a `MaterializedArray`, and seals the
+/// mutation surface so the wrapped module cannot be modified through this
+/// reference.
+///
+/// This also sets ``Module/training`` to `false`.
+///
+/// Note: only parameters that can be replaced, e.g. are wrapped with `@ParameterInfo`,
+/// will actually be updated to the `MaterializedArray` type.  All others will
+/// be evaluated and sealed, even if they do not have the type
+/// indicating that fact.
+///
+/// ## Construction and the `consuming` contract
+///
+/// The initializer takes its base module as `consuming`:
+///
+/// ```swift
+/// let lm = MaterializedModule(Linear(10, 10))
+/// ```
+///
+/// `consuming` expresses intent — **the caller must not retain or use the
+/// original `base` reference after passing it in**.  Because `Module` is a
+/// reference type, Swift cannot enforce this for you: a caller who keeps a
+/// reference and later mutates it will violate the `Sendable` invariant
+/// that `MaterializedModule` relies on.  As a runtime backstop, the base
+/// module and all of its descendants are sealed during initialization —
+/// any subsequent call to a mutating API (`update(parameters:)`,
+/// `update(modules:)`, `apply(...)`, `freeze`/`unfreeze`, `train`) on the
+/// retained reference will trap with `fatalError`.  In other words,
+/// `Sendable` here is a contract you can follow rather than a guarantee
+/// the compiler proves, but a violation fails loudly rather than silently.
+/// The recommended pattern is to construct the base module inline at the
+/// `MaterializedModule(...)` call site, as shown above, so no other
+/// reference exists.
+///
+/// ## What is sealed
+///
+/// `MaterializedModule` is not a `Module` and does not provide access to the
+/// wrapped `Module`.  No access to bare `MLXArray` is provided, though
+/// for purposes of introspection the ``parameters()`` can give
+/// `MaterializedArray`.  There are also properties like ``parameterNBytes``
+/// and ``logicalParameterCount`` for callers that just need size information.
+///
+/// The interior `Module` instances are still mutable objects, but they are
+/// not reachable through the public API.  They are flagged so that the
+/// `Module` mutation APIs trap.  Parameters that are held as bare
+/// `let x: MLXArray` cannot be replaced with `MaterializedArray`, but they
+/// are sealed in place, so in-place mutation of them (e.g. `x[0] = 0`) also
+/// traps.  Prefer `@ParameterInfo var x: MLXArray` so that the type reflects
+/// this.
+///
+/// Note that it is the `MaterializedModule` that is `Sendable` -- the
+/// interior `Module` instances are not.  Inference should not mutate
+/// module state (parameters in particular).  If a module does keep state
+/// that changes during a forward pass, e.g. a lazily built cache, it is
+/// responsible for synchronizing access to it (e.g. with a lock) so that
+/// concurrent calls through the `MaterializedModule` are safe.
+///
+/// ## Calling the wrapped module
+///
+/// `MaterializedModule` is not a `Module` subclass and cannot conform
+/// to protocols that have `Module` requirements.  Users may want to split their
+/// protocols into requirements for performing inference (which
+/// `MaterializedModule` is meant for) and for doing training where `Module`
+/// is required and this type is not appropriate (immutable).
+///
+/// `MaterializedModule` does not itself know how to invoke `base`; that is
+/// added per-layer-shape via an extension that constrains `LayerType`.
+/// For example, every ``UnaryLayer`` already supports being called with a
+/// single `MLXArray`, so the package ships:
+///
+/// ```swift
+/// extension MaterializedModule where LayerType: UnaryLayer {
+///     public func callAsFunction(_ x: MLXArray) -> MLXArray {
+///         _base(x)
+///     }
+/// }
+/// ```
+///
+/// A pattern that is used in `mlx-swift-lm` works like this.
+///
+/// The model is defined in terms of a protocol.  All language models (LLMs and VLMs)
+/// implement this protocol.  Note: in reality there are more methods and
+/// a `BaseLanguageModel` that lifts out methods useful to Embedding models.
+/// For this illustration it is simplified as those details are irrelevant to the technique.
+///
+/// ```swift
+/// public protocol LanguageModel {
+///     func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray
+/// }
+/// ```
+///
+/// Then an extension to `MaterializedModule` is created where the
+/// interior layer is of the given type and `MaterializedModule` is made
+/// to conform to the type:
+///
+/// ```swift
+/// // this allows access to _base
+/// @_spi(MaterializedModule) import MLXNN
+///
+/// extension MaterializedModule: LanguageModel where LayerType: LanguageModel {
+///
+///     public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray {
+///         _base.callAsFunction(inputs, cache: cache)
+///     }
+/// }
+/// ```
+///
+/// Finally there is a specific way that types can consume this.  Note the
+/// use of `any` and `some`.
+///
+/// ```swift
+/// public struct ModelContext: Sendable {
+///     public var model: any LanguageModel & Sendable
+///
+///     public init(
+///     model: some LanguageModel, ...
+///     ) {
+///         self.model = MaterializedModule(model)
+///         ...
+///     }
+/// }
+/// ```
+///
+/// ### Design Notes
+///
+/// Why use extensions?  Why not subclass?  Certainly it could, but it does require that the type you
+/// are specializing to is a subtype of `Module`.  You can't do this:
+///
+/// ```swift
+/// class MaterializedUnaryLayer: MaterializedModule<UnaryLayer> { ... }
+/// ```
+///
+/// because `UnaryLayer` is not a concrete type, it is a protocol.
+///
+/// The class is `open` so you could give it a try.
+open class MaterializedModule<LayerType: Module>: IndentedDescription, @unchecked Sendable {
+
+    /// Usable by extensions to implement `callAsFunction()`
+    @_spi(MaterializedModule)
+    public let _base: LayerType
+
+    /// Sum of all the `nbytes` of the parameters in the encapsulated model.
+    public let parameterNBytes: Int
+
+    /// Number of *logical* parameters in the encapsulated model, see
+    /// ``Module/logicalParameterCount``.
+    ///
+    /// For quantized layers this is the count of the equivalent unquantized
+    /// layer.  This is intended for display (e.g. "8B parameters") and is best
+    /// effort -- use ``parameterNBytes`` for the physical size.
+    public let logicalParameterCount: Int
+
+    public init(_ base: consuming LayerType) {
+        self._base = base
+
+        if !self._base._isSealed {
+            self._base.train(false)
+            self._base.materialize()
+
+            // seal the consumed base so that any retained reference held by a
+            // caller (despite the `consuming` contract) traps on mutation
+            // rather than silently violating Sendable
+            self._base._seal()
+        } else {
+            // it is already sealed.  could be a fatalError -- it suggests
+            // that the module was retained past encapsulation in a MaterializedModule,
+            // but is ultimately harmless.
+        }
+
+        parameterNBytes = self._base.parameters().reduce(0) { $0 + $1.nbytes }
+        logicalParameterCount = self._base.logicalParameterCount
+    }
+
+    /// Return a `NestedDictionary<String, MaterializedArray>` for all parameters in the
+    /// model (all layers).
+    public func parameters() -> NestedDictionary<String, MaterializedArray> {
+        // Note: the enclosed Module will have had all its parameters evaluated
+        // in init, but because it can't replace all parameters they may be typed
+        // as MLXArray.
+        _base.parameters().mapValues { $0.materialized() }
+    }
+
+    public func description(indent: Int) -> String {
+        _base.description(indent: indent)
+    }
+}
+
+extension MaterializedModule where LayerType: UnaryLayer {
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        _base(x)
+    }
+}
diff --git a/Source/MLXNN/Module.swift b/Source/MLXNN/Module.swift
index 39407c263..a2c0a0305 100644
--- a/Source/MLXNN/Module.swift
+++ b/Source/MLXNN/Module.swift
@@ -1,7 +1,7 @@
 // Copyright © 2024 Apple Inc.
 
 import Foundation
-import MLX
+@_spi(Sealing) import MLX
 
 /// `NestedDictionary` structure of `MLXArray`
 public typealias ModuleParameters = NestedDictionary<String, MLXArray>
@@ -109,6 +109,13 @@ open class Module {
     private var _items: ModuleItems?
     private var _setters: [String: TypeErasedSetter]?
 
+    /// Sealed flag.  When `true`, all mutation entry points
+    /// (parameter/module updates, freeze/unfreeze, train) trap.  Set
+    /// recursively by ``_seal()`` after the module's parameters
+    /// have been replaced with `MaterializedArray` values (or sealed in
+    /// place) inside ``MaterializedModule``.
+    var _isSealed = false
+
     /// Initializes the module.
     public init() {
     }
@@ -123,6 +130,8 @@ open class Module {
 
                 if let (_, _, setter) = isModuleInfo(c.value) {
                     setters[key] = setter
+                } else if let provider = c.value as? TypeErasedSetterProvider {
+                    setters[key] = provider.typeErasedSetter()
                 }
             }
 
@@ -185,6 +194,49 @@ open class Module {
         }
     }
 
+    /// Number of *logical* parameters at and beneath this module.
+    ///
+    /// This is intended for display, e.g. to describe the size of a loaded model
+    /// ("8B parameters").  It is not a measure of memory use -- for the physical
+    /// size sum the `nbytes` of ``parameters()``.
+    ///
+    /// For quantized layers (modules conforming to ``Quantized``) this reports the
+    /// count of the equivalent unquantized layer: the packed `weight` is counted as
+    /// `weight.size * 32 / bits` and the quantization metadata (`scales`, `biases`,
+    /// `global_scale`) is excluded.  Any other parameters (e.g. `bias` or parameters
+    /// added by a subclass, such as LoRA adapters) are counted as-is.
+    ///
+    /// This is best effort: it relies on the conventional parameter names above
+    /// and on quantized weights being packed into `uint32`.  Layers with unusual
+    /// storage may override it.
+    open var logicalParameterCount: Int {
+        let items = items()
+        var count = items.reduce(0) {
+            switch $1 {
+            case .parameters(let p):
+                $0 + p.size
+            case .module(let m):
+                $0 + m.logicalParameterCount
+            default:
+                $0
+            }
+        }
+
+        if let q = self as? Quantized, case .parameters(let w)? = items[unwrapping: "weight"] {
+            // replace the packed weight count with the logical (unpacked) count
+            count += w.size * 32 / q.bits - w.size
+
+            // quantization metadata is not part of the logical count
+            for key in ["scales", "biases", "global_scale"] {
+                if case .parameters(let p)? = items[unwrapping: key] {
+                    count -= p.size
+                }
+            }
+        }
+
+        return count
+    }
+
     /// Recursively filter and map the contents of the module and its children and produce a `NestedDictionary`
     /// with the results.
     ///
@@ -456,7 +508,19 @@ open class Module {
         parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [],
         modulePath: [String] = []
     ) throws -> Self {
+        _checkNotSealed()
+        try update(parameters: parameters, verify: verify, path: path, modulePath: modulePath) {
+            m, k, a, v in
+            a._updateInternal(v)
+        }
+        return self
+    }
 
+    func update(
+        parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [],
+        modulePath: [String] = [],
+        mutate: (Module, String, MLXArray, MLXArray) -> Void
+    ) throws {
         let modulePath = modulePath + [describeType(self)]
 
         func apply(
@@ -478,7 +542,7 @@ open class Module {
                         path: path, modules: modulePath, expectedShape: p.shape,
                         actualShape: newArray.shape)
                 }
-                p._updateInternal(newArray)
+                mutate(self, key, p, newArray)
 
             case (.value(.parameters), .none):
                 if Self.parameterIsValid(key) {
@@ -524,12 +588,12 @@ open class Module {
             case (.value(.module(let module)), .dictionary(let values)):
                 try module.update(
                     parameters: NestedDictionary(values: values), verify: verify, path: path,
-                    modulePath: modulePath)
+                    modulePath: modulePath, mutate: mutate)
 
             case (.value(.module(let module)), .none):
                 try module.update(
                     parameters: NestedDictionary(), verify: verify, path: path,
-                    modulePath: modulePath)
+                    modulePath: modulePath, mutate: mutate)
 
             case (.none, .none), (.value(.none), .none), (.value(.other(_)), .none):
                 break
@@ -555,8 +619,63 @@ open class Module {
             throw UpdateError.unhandledKeys(
                 path: path, modules: modulePath, keys: processed.sorted())
         }
+    }
 
-        return self
+    /// Recursively seal this module and all of its descendants so that any
+    /// further mutation (parameter updates, module replacement, freeze /
+    /// unfreeze, train mode) traps with `fatalError`.
+    ///
+    /// Called from ``MaterializedModule`` after the consumed base has been
+    /// materialized.  This is the runtime backstop for the `consuming`
+    /// ownership contract: even if a caller retains a stale reference and
+    /// tries to mutate the wrapped module, the call will trap rather than
+    /// silently violate the `Sendable` invariant.
+    func _seal() {
+        visit { _, m in
+            m._isSealed = true
+        }
+    }
+
+    private func _checkNotSealed(_ caller: StaticString = #function) {
+        if _isSealed {
+            fatalError(
+                "\(describeType(self)).\(caller): module has been sealed against "
+                    + "mutation by MaterializedModule.  The original reference "
+                    + "must not be retained or used after the module is consumed.")
+        }
+    }
+
+    func materialize() {
+        // bulk eval the parameters
+        eval(self.parameters())
+
+        // now convert to MaterializedArray (where possible)
+        let newParameters = filterMap(
+            filter: { _, _, _ in true },
+            map: Self.mapParameters(map: { $0.materialized() as MLXArray }))
+
+        // not verifying and setting with same value -- any
+        // errors are programming errors
+        try! update(parameters: newParameters, verify: .none) { m, k, a, v in
+            if let setter = m._setters?[k] {
+                do {
+                    // use the setter to replace the array
+                    try setter.update(v)
+                } catch {
+                    a._seal()
+                }
+            } else {
+                // can't replace it (e.g. `let`) -- it is already evaluated,
+                // so seal it in place.  `_seal()` is idempotent, so an array
+                // shared between several parameters (tied weights) is fine.
+                a._seal()
+            }
+        }
+
+        // some of the properties were updated so rebuild the properties cache
+        visit { key, m in
+            m.buildCaches()
+        }
     }
 
     /// Called from ``update(parameters:verify:path:modulePath:)`` if a required parameter
@@ -593,7 +712,8 @@ open class Module {
         filter: (Module, String, ModuleItem) -> Bool = Module.filterValidParameters,
         map: @escaping (MLXArray) -> MLXArray
     ) -> Self {
-        update(parameters: filterMap(filter: filter, map: Self.mapParameters(map: map)))
+        _checkNotSealed()
+        return update(parameters: filterMap(filter: filter, map: Self.mapParameters(map: map)))
     }
 
     /// A non-throwing version of ``update(modules:verify:path:modulePath:)``.
@@ -654,6 +774,7 @@ open class Module {
         modules: ModuleChildren, verify: VerifyUpdate, path: [String] = [],
         modulePath: [String] = []
     ) throws -> Self {
+        _checkNotSealed()
 
         let modulePath = modulePath + [describeType(self)]
 
@@ -798,13 +919,14 @@ open class Module {
     ///   - key: module key, see ``ModuleInfo``
     ///   - value: the replacement module
     open func updateModule(key: String, _ value: Any) throws {
+        _checkNotSealed()
         if _setters == nil {
             buildCaches()
         }
 
         if let setter = _setters?[key] {
             do {
-                try setter.updateModule(value)
+                try setter.update(value)
             } catch {
                 throw UpdateError.needModuleInfo(
                     "Unable to set modules for \(describeType(self)).\(key) -- maybe type mismatch: \(describeType(value)), \(error)"
@@ -912,6 +1034,7 @@ open class Module {
     /// - ``freeze(recursive:keys:)``
     /// - ``unfreeze(recursive:keys:strict:)``
     open func freeze(recursive: Bool = true, keys: [String]? = nil, strict: Bool = false) throws {
+        _checkNotSealed()
         let visitor = freezeVisitor(keys: keys, strict: strict) {
             $0._noGrad.formUnion($1)
             $0.didSetNoGrad($0._noGrad)
@@ -950,6 +1073,7 @@ open class Module {
     /// - ``Module/freeze(recursive:keys:)``
     /// - ``Module/unfreeze(recursive:keys:strict:)``
     open func unfreeze(recursive: Bool = true, keys: [String]? = nil, strict: Bool = false) throws {
+        _checkNotSealed()
         let visitor = freezeVisitor(keys: keys, strict: strict) {
             $0._noGrad.subtract($1)
             $0.didSetNoGrad($0._noGrad)
@@ -990,6 +1114,7 @@ open class Module {
     /// - ``training``
     /// - ``didSetTrain(_:)``
     public func train(_ mode: Bool = true) {
+        _checkNotSealed()
         visit(modules: {
             $1.training = mode
             $1.didSetTrain(mode)
@@ -1407,7 +1532,7 @@ public enum ModuleValue {
 /// ### See Also
 /// - <doc:custom-layers>
 /// - ``ModuleInfo``
-@propertyWrapper public class ParameterInfo<T> {
+@propertyWrapper public class ParameterInfo<T>: TypeErasedSetterProvider {
     var value: T?
     let key: String?
 
@@ -1453,11 +1578,48 @@ public enum ModuleValue {
 
         // cannot check via unwapProperty -- see wrappedValue.set
     }
+
+    // See also ModuleInfo.Setter
+    struct Setter: TypeErasedSetter {
+        unowned var info: ParameterInfo<T>
+
+        func update(_ value: Any) throws {
+            if let value = value as? T {
+                info.value = value
+            } else if let value = value as? [MLXArray] {
+                // try to recast as a tuple, e.g.
+                // @ParameterInfo var x: (MLXArray, MLXArray)
+
+                if value.count == 2, let values = (value[0], value[1]) as? T {
+                    info.value = values
+                } else if value.count == 3, let values = (value[0], value[1], value[2]) as? T {
+                    info.value = values
+                } else if value.count == 4,
+                    let values = (value[0], value[1], value[2], value[3]) as? T
+                {
+                    info.value = values
+                } else if value.count == 5,
+                    let values = (value[0], value[1], value[2], value[3], value[4]) as? T
+                {
+                    info.value = values
+                } else {
+                    throw UpdateError.unableToCast(String(describing: T.self))
+                }
+            } else {
+                throw UpdateError.unableToCast(String(describing: T.self))
+            }
+        }
+    }
+
+    fileprivate func typeErasedSetter() -> TypeErasedSetter {
+        Setter(info: self)
+    }
+
 }
 
 /// Helper protocol for writing back through ``ModuleInfo``, e.g. via ``Module/update(modules:)``
 private protocol TypeErasedSetter {
-    func updateModule(_ value: Any) throws
+    func update(_ value: Any) throws
 }
 
 private protocol TypeErasedSetterProvider {
@@ -1565,10 +1727,11 @@ private protocol TypeErasedSetterProvider {
         }
     }
 
+    // See also ParameterInfo.Setter
     struct Setter: TypeErasedSetter {
         unowned var info: ModuleInfo<T>
 
-        func updateModule(_ value: Any) throws {
+        func update(_ value: Any) throws {
             if let value = value as? T {
                 info.module = value
             } else if let value = value as? [Module] {
@@ -1584,7 +1747,7 @@ private protocol TypeErasedSetterProvider {
                 {
                     info.module = values
                 } else if value.count == 5,
-                    let values = (value[0], value[1], value[2], value[4], value[5]) as? T
+                    let values = (value[0], value[1], value[2], value[3], value[4]) as? T
                 {
                     info.module = values
                 } else {
diff --git a/Source/MLXNN/Quantized.swift b/Source/MLXNN/Quantized.swift
index b1e429794..840545580 100644
--- a/Source/MLXNN/Quantized.swift
+++ b/Source/MLXNN/Quantized.swift
@@ -165,8 +165,8 @@ open class QuantizedEmbedding: Embedding, Quantized {
     public let bits: Int
 
     public let mode: QuantizationMode
-    public let scales: MLXArray
-    public let biases: MLXArray?
+    @ParameterInfo public private(set) var scales: MLXArray
+    @ParameterInfo public private(set) var biases: MLXArray?
     @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray?
 
     open override var shape: (Int, Int) {
@@ -204,8 +204,8 @@ open class QuantizedEmbedding: Embedding, Quantized {
         let (quantizedWeight, scales, biases) = MLX.quantized(
             weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale)
 
-        self.scales = scales
-        self.biases = biases
+        self._scales.wrappedValue = scales
+        self._biases.wrappedValue = biases
 
         super.init(weight: quantizedWeight)
 
@@ -222,8 +222,8 @@ open class QuantizedEmbedding: Embedding, Quantized {
         self.groupSize = groupSize
         self.bits = bits
         self.mode = mode
-        self.scales = scales
-        self.biases = biases
+        self._scales.wrappedValue = scales
+        self._biases.wrappedValue = biases
         self._globalScale.wrappedValue = globalScale
         super.init(weight: weight)
     }
@@ -269,8 +269,8 @@ open class QuantizedLinear: Linear, Quantized {
     public let bits: Int
 
     public let mode: QuantizationMode
-    public let scales: MLXArray
-    public let biases: MLXArray?
+    @ParameterInfo public private(set) var scales: MLXArray
+    @ParameterInfo public private(set) var biases: MLXArray?
     @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray?
 
     open override var shape: (Int, Int) {
@@ -332,8 +332,8 @@ open class QuantizedLinear: Linear, Quantized {
         let (quantizedWeight, scales, biases) = MLX.quantized(
             weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale)
 
-        self.scales = scales
-        self.biases = biases
+        self._scales.wrappedValue = scales
+        self._biases.wrappedValue = biases
 
         super.init(weight: quantizedWeight, bias: bias)
 
@@ -353,12 +353,19 @@ open class QuantizedLinear: Linear, Quantized {
         self.groupSize = groupSize
         self.bits = bits
         self.mode = mode
-        self.scales = scales
-        self.biases = biases
+        self._scales.wrappedValue = scales
+        self._biases.wrappedValue = biases
         self._globalScale.wrappedValue = globalScale
         super.init(weight: weight, bias: bias)
     }
 
+    open override func describeExtra(_ indent: Int) -> String {
+        let outputDimensions = scales.dim(0)
+        let inputDimensions = scales.dim(1) * groupSize
+        return
+            "(inputDimensions=\(inputDimensions), outputDimensions=\(outputDimensions), bias=\(self.bias != nil))"
+    }
+
     public override func unfreeze(
         recursive: Bool = true, keys: [String]? = nil, strict: Bool = false
     ) throws {
diff --git a/Tests/MLXTests/LogicalParameterCountTests.swift b/Tests/MLXTests/LogicalParameterCountTests.swift
new file mode 100644
index 000000000..7ee366e4c
--- /dev/null
+++ b/Tests/MLXTests/LogicalParameterCountTests.swift
@@ -0,0 +1,132 @@
+// Copyright © 2026 Apple Inc.
+
+import Foundation
+import MLX
+import MLXNN
+import Testing
+
+@Suite
+struct LogicalParameterCountTests {
+
+    // MARK: - Test layers
+
+    /// A `QuantizedLinear` subclass with an extra parameter, e.g. shaped like
+    /// a QLoRA adapter.  The extra parameter must be counted.
+    final class QuantizedLinearWithExtra: QuantizedLinear {
+        let extra: MLXArray
+
+        init(_ inputDimensions: Int, _ outputDimensions: Int, extraRows: Int) {
+            self.extra = MLXArray.zeros([extraRows, inputDimensions])
+            super.init(
+                weight: MLXRandom.uniform(0 ..< 1, [outputDimensions, inputDimensions]),
+                bias: nil)
+        }
+    }
+
+    /// A ``Quantized`` layer that does not derive from `QuantizedLinear`, similar
+    /// to `QuantizedSwitchLinear` in mlx-swift-lm (a stack of expert weights).
+    final class QuantizedExperts: Module, Quantized {
+        let groupSize: Int
+        let bits: Int
+        let mode: QuantizationMode
+
+        let weight: MLXArray
+        let scales: MLXArray
+        let biases: MLXArray?
+        let bias: MLXArray
+
+        init(experts: Int, inputDimensions: Int, outputDimensions: Int, bits: Int = 4) {
+            self.groupSize = 64
+            self.bits = bits
+            self.mode = .affine
+
+            let w = MLXRandom.uniform(0 ..< 1, [experts, outputDimensions, inputDimensions])
+            let (wq, scales, biases) = MLX.quantized(w, groupSize: groupSize, bits: bits)
+            self.weight = wq
+            self.scales = scales
+            self.biases = biases
+            self.bias = MLXArray.zeros([experts, outputDimensions])
+        }
+    }
+
+    final class Model: Module {
+        @ModuleInfo var embedding: Embedding
+        @ModuleInfo var layers: [Linear]
+        @ModuleInfo var head: Linear
+
+        override init() {
+            self.embedding = Embedding(embeddingCount: 100, dimensions: 128)
+            self.layers = [Linear(128, 256), Linear(256, 128, bias: false)]
+            self.head = Linear(128, 100)
+        }
+    }
+
+    // MARK: - Tests
+
+    @Test func testLinear() {
+        #expect(Linear(128, 32).logicalParameterCount == 128 * 32 + 32)
+        #expect(Linear(128, 32, bias: false).logicalParameterCount == 128 * 32)
+    }
+
+    @Test(arguments: [
+        (4, QuantizationMode.affine, 64),
+        (8, QuantizationMode.affine, 64),
+        (3, QuantizationMode.affine, 64),
+        (4, QuantizationMode.mxfp4, 32),
+    ])
+    func testQuantizedLinear(bits: Int, mode: QuantizationMode, groupSize: Int) {
+        let withBias = QuantizedLinear(
+            256, 64, bias: true, groupSize: groupSize, bits: bits, mode: mode)
+        #expect(withBias.logicalParameterCount == 256 * 64 + 64)
+
+        let noBias = QuantizedLinear(
+            256, 64, bias: false, groupSize: groupSize, bits: bits, mode: mode)
+        #expect(noBias.logicalParameterCount == 256 * 64)
+    }
+
+    @Test func testQuantizedEmbedding() {
+        let e = QuantizedEmbedding(embeddingCount: 100, dimensions: 128)
+        #expect(e.logicalParameterCount == 100 * 128)
+    }
+
+    @Test func testQuantizedSubclassExtraParameters() {
+        let q = QuantizedLinearWithExtra(256, 64, extraRows: 8)
+        #expect(q.logicalParameterCount == 256 * 64 + 8 * 256)
+    }
+
+    @Test func testCustomQuantizedLayer() {
+        let q = QuantizedExperts(experts: 4, inputDimensions: 256, outputDimensions: 64)
+        #expect(q.logicalParameterCount == 4 * 256 * 64 + 4 * 64)
+    }
+
+    @Test func testNested() {
+        let model = Model()
+        let expected =
+            100 * 128  // embedding
+            + 128 * 256 + 256  // layers[0]
+            + 256 * 128  // layers[1]
+            + 128 * 100 + 100  // head
+        #expect(model.logicalParameterCount == expected)
+    }
+
+    @Test func testQuantizePreservesCount() {
+        // the logical count should not change when a model is quantized
+        let model = Model()
+        let before = model.logicalParameterCount
+
+        quantize(model: model, groupSize: 64, bits: 4, mode: .affine)
+        #expect(model.head is QuantizedLinear)
+        #expect(model.embedding is QuantizedEmbedding)
+
+        #expect(model.logicalParameterCount == before)
+    }
+
+    @Test func testMaterializedModule() {
+        let model = Model()
+        quantize(model: model, groupSize: 64, bits: 4, mode: .affine)
+        let expected = model.logicalParameterCount
+
+        let mm = MaterializedModule(model)
+        #expect(mm.logicalParameterCount == expected)
+    }
+}
diff --git a/Tests/MLXTests/MLXArray+OpsTests.swift b/Tests/MLXTests/MLXArray+OpsTests.swift
index 2c2cd3a58..684e10397 100644
--- a/Tests/MLXTests/MLXArray+OpsTests.swift
+++ b/Tests/MLXTests/MLXArray+OpsTests.swift
@@ -15,7 +15,7 @@ class MLXArrayOpsTests: XCTestCase {
 
     func testArithmeticSimple() {
         var a = MLXArray([1, 2, 3])
-        var b = MLXArray(converting: [-5.0, 37.5, 4])
+        let b = MLXArray(converting: [-5.0, 37.5, 4])
 
         // example of an expression -- the - 1 is using the 1 as ExpressibleByIntegerLiteral
         let r = a + b - 1
diff --git a/Tests/MLXTests/MaterializedHazardTests.swift b/Tests/MLXTests/MaterializedHazardTests.swift
new file mode 100644
index 000000000..5bc940211
--- /dev/null
+++ b/Tests/MLXTests/MaterializedHazardTests.swift
@@ -0,0 +1,597 @@
+// Copyright © 2026 Apple Inc.
+
+import Foundation
+@_spi(MaterializedModule) import MLXNN
+import Testing
+
+@testable import MLX
+
+// Tests for ways a MaterializedArray / MaterializedModule can end up
+// mutable or not actually materialized.
+//
+// These assert the *desired* invariants, so some of them are expected to
+// FAIL until the corresponding fix lands -- they demonstrate the hazard.
+// Each test is annotated with what it is expected to observe.
+//
+// 1. `materialize()` inside `grad`/`vjp` would return a MaterializedArray
+//    that still carries the traced graph (primitive + inputs), and that graph
+//    is later mutated (detached) by an ordinary eval.  It traps instead.
+// 2a. `compile(inputs:outputs:)` swapped the parameters of a sealed module
+//     for tracers when they are plain `let` (not `@ParameterInfo`) and so were
+//     not replaced with MaterializedArray.  They are now sealed in place and
+//     compile treats sealed arrays as constants.
+// 2b. `Linear.weight` / `Linear.bias` have public setters (intentionally --
+//     outside a MaterializedModule a Linear is mutable), which bypass the
+//     sealed-module checks.
+// 3. `compile(inputs:outputs:)` bakes sealed state into the trace as
+//    constants, but the compile cache does not key on them -- if the state is
+//    rebound to different arrays (a different module, a replaced child, a new
+//    MaterializedArray) a later call could reuse the stale trace.
+//    `CompiledFunction` now checks the identity of the sealed arrays and
+//    retraces when they change.
+// 4. The `@Sendable` closure APIs (`compile(shapeless:_:)`, `vmapSendable`,
+//    `CustomFunction`): these can capture MaterializedArray /
+//    MaterializedModule (they are Sendable), and `materialize` must still be
+//    rejected inside them while they are being traced.
+
+// MARK: - Support
+
+/// True if `array` still references a computation graph, i.e. it has a
+/// primitive (and so inputs/siblings) rather than being a detached leaf.
+///
+/// `graphDescription` prints `Inputs:` and `Outputs:` lines followed by one
+/// line per primitive in the graph -- a detached, evaluated array has none.
+private func hasGraph(_ array: MLXArray) -> Bool {
+    graphDescription([array])
+        .split(separator: "\n")
+        .dropFirst(2)
+        .contains { !$0.isEmpty }
+}
+
+/// True if `a` and `b` refer to the same underlying `mlx::core::array` node
+/// (the node ids match), regardless of which `MLXArray` wrapper holds them.
+private func sameNode(_ a: MLXArray, _ b: MLXArray) -> Bool {
+    let namer = NodeNamer()
+    namer.setName("target", for: a)
+    return namer.name(for: b) == "target"
+}
+
+// MARK: - Bodies for exit tests (must not capture state)
+
+/// (1) `materialize` of a traced value inside `grad`.
+private func materializeInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        materialize(x * x).sum()
+    }
+    eval(g(MLXArray(Float(3))))
+}
+
+/// (1) batch `materialize` of traced values inside `vjp`.
+private func materializeArraysInsideVJP() {
+    let (_, ct) = vjp(
+        { (inputs: [MLXArray]) -> [MLXArray] in
+            let m = materialize([inputs[0] * 2, inputs[0] + 1])
+            return [m[0] * m[1]]
+        }, primals: [MLXArray(Float(3))], cotangents: [MLXArray(Float(1))])
+    eval(ct)
+}
+
+/// (1) `materialized()` of a value that does not depend on the traced input,
+/// inside `grad`.
+private func materializeConstantInsideGrad() {
+    let c = MLXArray(Float(2))
+    let g = grad { (x: MLXArray) -> MLXArray in
+        (x * (c + 1).materialized()).sum()
+    }
+    eval(g(MLXArray(Float(3))))
+}
+
+/// (1) constructing a MaterializedModule inside `grad`.
+private func materializedModuleInsideGrad() {
+    let g = grad { (x: MLXArray) -> MLXArray in
+        let mm = MaterializedModule(Linear(1, 1))
+        return mm(x.reshaped(1, 1)).sum()
+    }
+    eval(g(MLXArray(Float(3))))
+}
+
+/// (1) `materialized()` of an existing MaterializedArray inside `grad`.
+private func rematerializeInsideGrad() {
+    let c = MLXArray(Float(2)).materialized()
+    let g = grad { (x: MLXArray) -> MLXArray in
+        (x * c.materialized()).sum()
+    }
+    #expect(g(MLXArray(Float(3))).item(Float.self) == 2)
+}
+
+/// (2a) A sealed module shared between a thread that is running
+/// `compile(inputs:outputs:)` on it and threads that are just running the
+/// forward pass.  If the compile trace swaps the (plain `let`) parameters for
+/// tracers, the forward threads can build graphs on tracer placeholders --
+/// which fail to evaluate (fatal error) or produce wrong results.  This is a
+/// stress test, so a pass is not proof of absence; it is most useful under
+/// the thread sanitizer.
+private func compileSealedModuleWhileShared() {
+    let mm = MaterializedModule(
+        Sequential(layers: RMSNorm(dimensions: 16), LayerNorm(dimensions: 16)))
+    let x = MLXRandom.normal([4, 16]).materialized()
+    let expected = mm(x).materialized()
+
+    DispatchQueue.concurrentPerform(iterations: 8) { worker in
+        for _ in 0 ..< 100 {
+            if worker == 0 {
+                // a fresh CompiledFunction each time so that it traces each time
+                let base = mm._base
+                let compiled = compile(inputs: [base], outputs: [base]) { (input: MLXArray) in
+                    mm(input)
+                }
+                eval(compiled(x))
+            } else {
+                let r = mm(x)
+                #expect(r.allClose(expected).item(Bool.self))
+            }
+        }
+    }
+}
+
+/// (2b) Assign `Linear.weight` on a sealed module.  The module is only
+/// reachable through `_base` (SPI) or a reference retained in violation of
+/// the `MaterializedModule` contract.
+private func writeWeightOfSealedLinear() {
+    let mm = MaterializedModule(Linear(4, 4))
+    mm._base.weight = MLXArray.zeros([4, 4])
+}
+
+/// (2b) Assign `Linear.bias` on a sealed module built with `bias: false`.
+private func writeNilBiasOfSealedLinear() {
+    let mm = MaterializedModule(Linear(4, 4, bias: false))
+    mm._base.bias = MLXArray.ones([4])
+}
+
+// MARK: - Tests
+
+@Suite("MaterializedArray / MaterializedModule hazards")
+struct MaterializedHazardTests {
+
+    // MARK: (1) materialize() inside a transform
+
+    /// Control: outside of any transform `materialize` produces a detached
+    /// leaf.  This also validates `hasGraph`.
+    @Test func materializeOutsideTransformIsDetached() {
+        let x = MLXArray(Float(3))
+        let y = x * x
+        #expect(hasGraph(y))
+
+        let m = materialize(y)
+        #expect(!hasGraph(m))
+    }
+
+    /// Traps.  Inside `grad` an eval of a traced value does not detach it, so
+    /// the result would still reference the trace (and be mutated by a later
+    /// eval) -- it cannot be a MaterializedArray.
+    @Test func materializeInsideGradTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeInsideGrad()
+        }
+    }
+
+    /// Traps.  Same via the batch overload and `vjp`.
+    @Test func materializeArraysInsideVJPTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeArraysInsideVJP()
+        }
+    }
+
+    /// Traps.  The check is "inside a transformation", not "is a tracer", so
+    /// this is rejected even though the array does not depend on the traced
+    /// input.
+    @Test func materializeConstantInsideGradTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeConstantInsideGrad()
+        }
+    }
+
+    /// Traps.  Constructing a MaterializedModule materializes its parameters.
+    @Test func materializedModuleInsideGradTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializedModuleInsideGrad()
+        }
+    }
+
+    /// Allowed: `materialized()` on an array that is already a
+    /// MaterializedArray returns it unchanged and does not trap.
+    @Test func rematerializeInsideGradIsAllowed() async {
+        await #expect(processExitsWith: .success) {
+            rematerializeInsideGrad()
+        }
+    }
+
+    // MARK: (2a) compile(inputs:outputs:) on a sealed module
+
+    /// Passes.  `RMSNorm.weight` and `LayerNorm.weight/bias` are plain `let`
+    /// and so stay `MLXArray` after sealing, but they are sealed in place.
+    /// `compile` must not swap them for tracers during the trace (which
+    /// would be visible to any other thread using the shared module) or write
+    /// the compiled outputs back into them afterwards.
+    @Test func compileInputsOutputsDoesNotSwapSealedParameters() {
+        let mm = MaterializedModule(
+            Sequential(layers: RMSNorm(dimensions: 16), LayerNorm(dimensions: 16)))
+        let base = mm._base
+
+        let state = base.innerState()
+
+        // precondition for the test: these are the parameters that could not
+        // be replaced with MaterializedArray
+        #expect(state.count == 3)
+        #expect(state.allSatisfy { !($0 is MaterializedArray) })
+
+        // independent handles on the original nodes
+        let originals = state.map { $0.copyContext() }
+
+        var traced = false
+        var swappedDuringTrace = [Bool]()
+
+        let compiled = compile(inputs: [base], outputs: [base]) { (x: MLXArray) in
+            traced = true
+            swappedDuringTrace = zip(base.innerState(), originals).map { !sameNode($0, $1) }
+            return mm(x)
+        }
+
+        let x = MLXRandom.normal([4, 16])
+        let expected = mm(x)
+        let result = compiled(x)
+        #expect(result.allClose(expected).item(Bool.self))
+
+        // make sure the trace actually ran so the checks below mean something
+        #expect(traced)
+        #expect(
+            swappedDuringTrace.allSatisfy { !$0 },
+            "sealed parameters were swapped for tracers during compile: \(swappedDuringTrace)")
+
+        let replacedAfterCall = zip(base.innerState(), originals).map { !sameNode($0, $1) }
+        #expect(
+            replacedAfterCall.allSatisfy { !$0 },
+            "sealed parameters were replaced by compile outputs: \(replacedAfterCall)")
+    }
+
+    /// Passes.  This was a race (crash or mismatch) when compile swapped the
+    /// sealed parameters, so it is most useful under TSan.
+    @Test func compileSealedModuleWhileSharedAcrossThreads() async {
+        await #expect(processExitsWith: .success) {
+            compileSealedModuleWhileShared()
+        }
+    }
+
+    // MARK: (2b) Linear.weight / Linear.bias setters
+
+    // The setters are public: outside a MaterializedModule a Linear is an
+    // ordinary mutable Module.  Writing them on a *sealed* module should trap,
+    // like every other mutation of a sealed module.
+
+    /// Passes -- the write traps, but in `ParameterInfo.wrappedValue.set`
+    /// ("please call update()", for any non-nil current value), not in
+    /// `_checkNotSealed()`.
+    @Test func writeWeightOfSealedLinearTraps() async {
+        await #expect(processExitsWith: .failure) {
+            writeWeightOfSealedLinear()
+        }
+    }
+
+    /// Passes -- `Linear.init` assigns `bias = nil`, which `ParameterInfo`
+    /// stores as `.some(nil)`, so this also hits the "please call update()"
+    /// trap.  (Only an optional `@ParameterInfo` that `init` never assigns
+    /// can be set later.)
+    @Test func writeNilBiasOfSealedLinearTraps() async {
+        await #expect(processExitsWith: .failure) {
+            writeNilBiasOfSealedLinear()
+        }
+    }
+}
+
+// MARK: - (3) compile bakes materialized state but the cache ignores the binding
+
+/// Holds a module that can be rebound -- the indirection that lets a
+/// compiled function see a different module than the one it was traced with.
+private final class ModuleHolder<M: UnaryLayer>: Updatable {
+    var model: M
+    init(_ model: M) { self.model = model }
+    func innerState() -> [MLXArray] { model.innerState() }
+}
+
+/// Holds a single array that can be rebound, e.g. state that is replaced
+/// each step.
+private final class ArrayHolder: Updatable {
+    var value: MLXArray
+    init(_ value: MLXArray) { self.value = value }
+    func innerState() -> [MLXArray] { [value] }
+}
+
+/// A plain module container whose child can be replaced via
+/// `update(modules:)`.
+private final class Parent: Module, UnaryLayer {
+    @ModuleInfo var child: Linear
+    init(_ child: Linear) { self._child.wrappedValue = child }
+    func callAsFunction(_ x: MLXArray) -> MLXArray { child(x) }
+}
+
+/// A `Linear` with every weight and bias element equal to `value`, so that
+/// different instances give clearly different results.  If `sealed`, it is
+/// the (sealed) base of a MaterializedModule: all of its parameters are
+/// `@ParameterInfo` so they are all MaterializedArray.
+private func makeLinear(_ value: Float, sealed: Bool) -> Linear {
+    let l = Linear(weight: MLXArray.ones([4, 4]) * value, bias: MLXArray.ones([4]) * value)
+    return sealed ? MaterializedModule(l)._base : l
+}
+
+/// `RMSNorm` (plain `let` weight, stays `MLXArray`) followed by `Linear`
+/// (`@ParameterInfo`, becomes MaterializedArray) -- a partly materialized
+/// module.  RMSNorm is first so the Linear scale shows in the output.
+private func partlySealed(_ value: Float) -> Sequential {
+    let model = Sequential(layers: RMSNorm(dimensions: 4), makeLinear(value, sealed: false))
+    return MaterializedModule(model)._base
+}
+
+/// A fixed (non-random) input for the compiled functions.
+private func compileInput() -> MLXArray {
+    MLXArray((0 ..< 8).map { Float($0) }, [2, 4])
+}
+
+@Suite("compile with rebound materialized state")
+struct CompileReboundMaterializedStateTests {
+
+    /// Control: plain modules rebound between calls are threaded through as
+    /// inputs and the compiled function follows the rebinding.
+    @Test func plainModuleRebound() {
+        let input = compileInput()
+        let holder = ModuleHolder(makeLinear(1, sealed: false))
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
+
+        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+
+        holder.model = makeLinear(2, sealed: false)
+        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+    }
+
+    /// Control: a plain array rebound between calls.
+    @Test func plainArrayRebound() {
+        let input = compileInput()
+        let holder = ArrayHolder(MLXArray(Float(1)))
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in x + holder.value }
+
+        #expect(c(input).allClose(input + 1).item(Bool.self))
+
+        holder.value = MLXArray(Float(2))
+        #expect(c(input).allClose(input + 2).item(Bool.self))
+    }
+
+    /// Sealed A -> sealed A2 (same shapes): both have zero threaded state
+    /// inputs, so without the identity check the cache would hit the trace
+    /// with A's weights baked in.
+    @Test func sealedModuleReboundToSealed() {
+        let input = compileInput()
+        let holder = ModuleHolder(makeLinear(1, sealed: true))
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
+
+        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+
+        holder.model = makeLinear(2, sealed: true)
+        #expect(
+            c(input).allClose(holder.model(input)).item(Bool.self),
+            "compiled function used the weights of the module it was traced with")
+    }
+
+    /// Sealed A (0 state inputs) -> plain B (2) retraces on the input count
+    /// alone, but sealed A2 (0 again) would match A's cache entry without the
+    /// identity check.
+    @Test func sealedPlainSealed() {
+        let input = compileInput()
+        let holder = ModuleHolder(makeLinear(1, sealed: true))
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
+
+        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+
+        holder.model = makeLinear(2, sealed: false)
+        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+
+        holder.model = makeLinear(3, sealed: true)
+        #expect(
+            c(input).allClose(holder.model(input)).item(Bool.self),
+            "compiled function reused the trace of an earlier sealed module")
+    }
+
+    /// Partly sealed A -> partly sealed B: the same number of threaded inputs
+    /// (the RMSNorm weight), so without the identity check the cache would hit
+    /// and combine B's RMSNorm weight with A's baked Linear weights.
+    @Test func partlySealedRebound() {
+        let input = compileInput()
+        let holder = ModuleHolder(partlySealed(1))
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
+
+        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+
+        holder.model = partlySealed(2)
+        #expect(
+            c(input).allClose(holder.model(input)).item(Bool.self),
+            "compiled function mixed baked weights from the traced module")
+    }
+
+    /// Same as `sealedModuleReboundToSealed` but through a
+    /// plain (unsealed) parent `Module` and `update(modules:)`, the usual way
+    /// to swap a layer.
+    @Test func sealedChildReplacedViaUpdateModules() {
+        let input = compileInput()
+        let parent = Parent(makeLinear(1, sealed: true))
+        let c = compile(inputs: [parent], outputs: [parent]) { (x: MLXArray) in parent(x) }
+
+        #expect(c(input).allClose(parent(input)).item(Bool.self))
+
+        var u = NestedDictionary<String, Module>()
+        u["child"] = .value(makeLinear(2, sealed: true))
+        parent.update(modules: u)
+
+        #expect(
+            c(input).allClose(parent(input)).item(Bool.self),
+            "compiled function used the weights of the child it was traced with")
+    }
+
+    /// State that is replaced by a new MaterializedArray each step: without
+    /// the identity check the first value is baked in and later ones are
+    /// ignored.  (This retraces on every rebinding -- state that changes every
+    /// step should be a plain array.)
+    @Test func materializedArrayRebound() {
+        let input = compileInput()
+        let holder = ArrayHolder(MLXArray(Float(1)).materialized())
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in x + holder.value }
+
+        #expect(c(input).allClose(input + 1).item(Bool.self))
+
+        holder.value = MLXArray(Float(2)).materialized()
+        #expect(
+            c(input).allClose(input + 2).item(Bool.self),
+            "compiled function used the MaterializedArray it was traced with")
+    }
+
+    /// The identity check must not retrace while the binding is stable, and
+    /// must retrace (once) when it changes.
+    @Test func sealedRetracesOnlyWhenRebound() {
+        let input = compileInput()
+        let holder = ModuleHolder(makeLinear(1, sealed: true))
+
+        var traces = 0
+        let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in
+            traces += 1
+            return holder.model(x)
+        }
+
+        for _ in 0 ..< 3 {
+            eval(c(input))
+        }
+        #expect(traces == 1)
+
+        holder.model = makeLinear(2, sealed: true)
+        for _ in 0 ..< 3 {
+            #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        }
+        #expect(traces == 2)
+    }
+}
+
+// MARK: - (4) @Sendable compile / vmap / CustomFunction
+
+/// (4) `materialize` inside a `@Sendable` compiled function.
+private func materializeInsideSendableCompile() {
+    let c: @Sendable (MLXArray) -> MLXArray = compile { x in
+        (x * 2).materialized()
+    }
+    eval(c(MLXArray(Float(3))))
+}
+
+/// (4) `materialize` inside a `vmapSendable` function.
+private func materializeInsideVmapSendable() {
+    let f = vmapSendable { (x: MLXArray) -> MLXArray in
+        (x * 2).materialized()
+    }
+    eval(f(MLXArray([Float(1), 2, 3])))
+}
+
+/// (4) `materialize` inside a `CustomFunction` forward that is run by `grad`.
+private func materializeInsideCustomForwardUnderGrad() {
+    let f = CustomFunction {
+        Forward { inputs in [(inputs[0] * 2).materialized()] }
+        VJP { _, cotangents in [cotangents[0] * 2] }
+    }
+    let g = grad { (x: MLXArray) -> MLXArray in f([x])[0].sum() }
+    eval(g(MLXArray(Float(3))))
+}
+
+/// (4) `materialize` inside a `CustomFunction` VJP (run by `grad`).
+private func materializeInsideCustomVJPUnderGrad() {
+    let f = CustomFunction {
+        Forward { inputs in [inputs[0] * 2] }
+        VJP { _, cotangents in [(cotangents[0] * 2).materialized()] }
+    }
+    let g = grad { (x: MLXArray) -> MLXArray in f([x])[0].sum() }
+    eval(g(MLXArray(Float(3))))
+}
+
+@Suite("@Sendable transforms with materialized state")
+struct SendableTransformsMaterializedTests {
+
+    /// A `@Sendable` compiled function may capture a MaterializedModule (its
+    /// parameters are sealed constants) and be called concurrently.
+    @Test func sendableCompileCapturingMaterializedModule() async {
+        let model = MaterializedModule(makeLinear(1, sealed: false))
+        let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in
+            model(x)
+        }
+
+        let input = compileInput().materialized()
+        let expected = model(input).materialized()
+
+        await withTaskGroup(of: Bool.self) { group in
+            for _ in 0 ..< 8 {
+                group.addTask {
+                    compiled(input).allClose(expected).item(Bool.self)
+                }
+            }
+            for await ok in group {
+                #expect(ok)
+            }
+        }
+    }
+
+    /// A `@Sendable` compiled function may capture a MaterializedArray.
+    @Test func sendableCompileCapturingMaterializedArray() {
+        let scale = MLXArray(Float(3)).materialized()
+        let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in
+            x * scale
+        }
+        let input = compileInput()
+        #expect(compiled(input).allClose(input * 3).item(Bool.self))
+    }
+
+    /// `CustomFunction` closures may capture a MaterializedArray, including
+    /// the VJP that is run by `grad`.
+    @Test func customFunctionCapturingMaterializedArray() {
+        let scale = MLXArray(Float(3)).materialized()
+        let f = CustomFunction {
+            Forward { inputs in [inputs[0] * scale] }
+            VJP { _, cotangents in [cotangents[0] * scale] }
+        }
+        let g = grad { (x: MLXArray) -> MLXArray in f([x])[0].sum() }
+        #expect(g(MLXArray(Float(2))).item(Float.self) == 3)
+    }
+
+    /// Allowed: a `CustomFunction` forward that is *not* being run by a
+    /// transformation is ordinary code, so `materialize` works.
+    @Test func materializeInsideCustomForwardOutsideTransform() {
+        let f = CustomFunction {
+            Forward { inputs in [(inputs[0] * 2).materialized()] }
+        }
+        #expect(f([MLXArray(Float(3))])[0].item(Float.self) == 6)
+    }
+
+    @Test func materializeInsideSendableCompileTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeInsideSendableCompile()
+        }
+    }
+
+    @Test func materializeInsideVmapSendableTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeInsideVmapSendable()
+        }
+    }
+
+    @Test func materializeInsideCustomForwardUnderGradTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeInsideCustomForwardUnderGrad()
+        }
+    }
+
+    @Test func materializeInsideCustomVJPUnderGradTraps() async {
+        await #expect(processExitsWith: .failure) {
+            materializeInsideCustomVJPUnderGrad()
+        }
+    }
+}
diff --git a/Tests/MLXTests/MaterializedTests.swift b/Tests/MLXTests/MaterializedTests.swift
new file mode 100644
index 000000000..cdc55f185
--- /dev/null
+++ b/Tests/MLXTests/MaterializedTests.swift
@@ -0,0 +1,270 @@
+// Copyright © 2024 Apple Inc.
+
+import Foundation
+import MLX
+@_spi(MaterializedModule) import MLXNN
+import Testing
+
+@Suite
+struct MaterializedTests {
+    @Test
+    func testMaterialize() async {
+        // a materialized array is Sendable
+        let x = MLXArray(10) + 5
+        let m = materialize(x)
+        let t = Task {
+            print(m + 3)
+        }
+        _ = await t.result
+    }
+
+    @Test
+    func testCompileMaterialized() async {
+        // compile manipulates the input arrays (tracers)
+        // make sure MaterializedArray doesn't run into problems here
+        func f(_ a: MLXArray, _ b: MLXArray) -> MLXArray {
+            square(a * b)
+        }
+
+        let c = compile(f)
+
+        let i1 = MLXRandom.normal([20, 20]).materialized()
+        let i2 = MLXRandom.normal([20, 20]).materialized()
+
+        let t = Task {
+            let s = sum(c(i1, i2))
+            let s2 = sum(f(i1, i2))
+            #expect(s.allClose(s2).item(Bool.self))
+            print(s, s2)
+        }
+        _ = await t.result
+    }
+
+    @Test
+    func testCompileInputsOutputsMaterialized() {
+        // A MaterializedArray can never change, so it has nothing for
+        // compile(inputs:outputs:) to observe or update -- it is simply
+        // captured as a constant when the function is traced, the same as
+        // any other value reachable from the closure that isn't part of
+        // `inputs`/`outputs`.  This is the documented
+        // `compile(inputs: [state], outputs: [state], f)` pattern (see
+        // compilation.md).
+        let state = MLXRandom.normal([4, 4]).materialized()
+        let compiled = compile(inputs: [state], outputs: [state]) { (x: MLXArray) in
+            x + state
+        }
+
+        let x = MLXRandom.normal([4, 4])
+        let r1 = compiled(x)
+        let r2 = x + state
+        #expect(r1.allClose(r2).item(Bool.self))
+
+        // same for a `Module` whose parameters were materialized (e.g. via
+        // `MaterializedModule`) -- this is the documented
+        // `compile(inputs: [model], outputs: [model])` pattern.
+        //
+        // Note: capturing the interior module is documented as forbidden but
+        // used here to test.  Do not copy this pattern!
+        let l = Linear(4, 4)
+        _ = MaterializedModule(l)
+
+        let compiledModel = compile(inputs: [l], outputs: [l]) { (x: MLXArray) in
+            l(x)
+        }
+
+        let input = MLXRandom.normal([4, 4])
+        let modelResult = compiledModel(input)
+        let expected = l(input)
+        #expect(modelResult.allClose(expected).item(Bool.self))
+    }
+
+    @Test
+    func testMaterializedLinear() async {
+        let l = Linear(10, 10)
+        let lm = MaterializedModule(l)
+
+        // this will have been materialized in the call
+        #expect(l.weight is MaterializedArray)
+
+        let t = Task {
+            let i = MLXRandom.normal([10, 10])
+            let r = lm(i)
+            print(sum(r))
+            print(lm)
+        }
+        _ = await t.result
+
+    }
+
+    @Test
+    func testMaterializedMultithreadedEval() async throws {
+        // Exercise concurrent evaluation: produce MaterializedArrays on the
+        // main task, fan them out to many child tasks, and have each task do
+        // its own work (creates new arrays, evals, reads scalars) at the same
+        // time.  This stresses the evalLock and the Sendable contract of
+        // MaterializedArray.
+
+        let taskCount = 16
+        let iterations = 8
+        let shape = [32, 32]
+
+        // shared inputs created up-front and materialized so they can cross
+        // task boundaries
+        let a = MLXRandom.normal(shape).materialized()
+        let b = MLXRandom.normal(shape).materialized()
+
+        // expected reference value computed serially on the main task
+        let expected = sum(square(a * b)).item(Float.self)
+
+        try await withThrowingTaskGroup(of: Float.self) { group in
+            for i in 0 ..< taskCount {
+                group.addTask {
+                    var last: Float = 0
+                    for _ in 0 ..< iterations {
+                        // mix the shared inputs with task-local arrays so that
+                        // each task is producing fresh graphs and evaluating
+                        // them concurrently
+                        let local = MLXRandom.normal(shape)
+                        let r = sum(square(a * b) + (local - local))
+                        last = r.item(Float.self)
+
+                        // also exercise materialize from inside a task
+                        let m = (a + MLXArray(Float(i))).materialized()
+                        _ = (m - MLXArray(Float(i))).sum().item(Float.self)
+                    }
+                    return last
+                }
+            }
+
+            for try await value in group {
+                #expect(abs(value - expected) < 1e-2)
+            }
+        }
+    }
+
+    @Test
+    func testMaterializedHighContention() async throws {
+        // High-contention variant: many tasks producing and consuming
+        // MaterializedArrays through a shared actor-protected pool.  Every
+        // task in a tight loop pulls two arrays from the pool, computes a
+        // new one, materializes it (forcing an eval), reads a scalar, and
+        // pushes the result back into the pool.  Lots of arrays flowing
+        // between tasks, lots of concurrent eval calls hammering the
+        // evalLock.
+
+        actor Pool {
+            var arrays: [MaterializedArray]
+            init(_ initial: [MaterializedArray]) { self.arrays = initial }
+
+            func take() -> MaterializedArray {
+                arrays.randomElement()!
+            }
+
+            func replace(_ a: MaterializedArray) {
+                arrays[Int.random(in: 0 ..< arrays.count)] = a
+            }
+
+            func snapshot() -> [MaterializedArray] {
+                arrays
+            }
+        }
+
+        let shape = [16, 16]
+        let poolSize = 16
+        let taskCount = 32
+        let iterations = 50
+
+        let initial = (0 ..< poolSize).map { _ in
+            MLXRandom.normal(shape).materialized()
+        }
+        let pool = Pool(initial)
+
+        // sum + count tracker so we can assert no task silently dropped work
+        actor Counter {
+            var n = 0
+            func bump() { n += 1 }
+            func value() -> Int { n }
+        }
+        let counter = Counter()
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            for t in 0 ..< taskCount {
+                group.addTask {
+                    for k in 0 ..< iterations {
+                        // pull two arrays from the shared pool — these may be
+                        // referenced concurrently by other tasks at the same time
+                        let a = await pool.take()
+                        let b = await pool.take()
+
+                        // build a new graph, materialize it (forces eval inside
+                        // the task), and read a scalar (forces another eval).
+                        // tanh keeps values bounded to [-1, 1] so the pool can
+                        // recycle results across iterations without blowing up.
+                        let mixin = MLXArray(Float(t * 31 + k) * 1e-3)
+                        let r = tanh((a * b) + (a - b) + mixin).materialized()
+                        let s = r.sum().item(Float.self)
+                        #expect(s.isFinite)
+
+                        // put the result back so other tasks see fresh arrays
+                        await pool.replace(r)
+                        await counter.bump()
+                    }
+                }
+            }
+            try await group.waitForAll()
+        }
+
+        #expect(await counter.value() == taskCount * iterations)
+
+        // every entry in the final pool should still be a valid, finite array
+        let final = await pool.snapshot()
+        #expect(final.count == poolSize)
+        for a in final {
+            #expect(a.shape == shape)
+            #expect(sum(a).item(Float.self).isFinite)
+        }
+    }
+
+    @Test func testMaterializedModule() {
+        let inner = Linear(10, 10)
+        let mm = MaterializedModule(inner)
+
+        #expect(mm.description.contains("Linear"))
+        #expect(mm.parameterNBytes == mm.parameters().reduce(0) { $0 + $1.nbytes })
+
+        // callable
+        let x = uniform(0 ..< 1, [10])
+        let _ = mm(x)
+    }
+
+    @Test func testMaterializedModuleQuantized() {
+        // scales and biases on QuantizedLinear are @ParameterInfo and
+        // should be replaced with MaterializedArray
+        let mm = MaterializedModule(QuantizedLinear(128, 32, bias: true))
+
+        #expect(mm._base.weight is MaterializedArray)
+        #expect(mm._base.bias is MaterializedArray)
+        #expect(mm._base.scales is MaterializedArray)
+        #expect(mm._base.biases is MaterializedArray)
+
+        let x = uniform(0 ..< 1, [128])
+        let _ = mm(x)
+    }
+
+    @Test func testMaterializedModuleTwice() {
+        let inner = Linear(10, 10)
+        let mm = MaterializedModule(inner)
+
+        // per contract with MaterializedModule, we can't retain the module
+        // after materializing it, but if we do, it should
+        let mm2 = MaterializedModule(inner)
+
+        #expect(mm.description.contains("Linear"))
+        #expect(mm.parameterNBytes == mm.parameters().reduce(0) { $0 + $1.nbytes })
+
+        // callable
+        let x = uniform(0 ..< 1, [10])
+        let _ = mm(x)
+        let _ = mm2(x)
+    }
+}
diff --git a/skills/mlx-swift/SKILL.md b/skills/mlx-swift/SKILL.md
index f3bc120f5..f0fc06555 100644
--- a/skills/mlx-swift/SKILL.md
+++ b/skills/mlx-swift/SKILL.md
@@ -40,6 +40,8 @@ Cmlx (C/C++ bindings, Metal GPU)
 | Purpose | File Path |
 |---------|-----------|
 | Core array | Source/MLX/MLXArray.swift |
+| Materialized (Sendable) array | Source/MLX/MaterializedArray.swift |
+| Materialized (Sendable) module | Source/MLXNN/MaterializedModule.swift |
 | Operations | Source/MLX/Ops.swift |
 | Transforms | Source/MLX/Transforms.swift |
 | Factory methods | Source/MLX/Factory.swift |
@@ -387,7 +389,7 @@ _ = await weightsTicket.end()
 
 ### DON'T
 
-- **Don't share MLXArrays across tasks**: MLXArray is NOT Sendable by design.
+- **Don't share MLXArrays across tasks**: MLXArray is NOT Sendable by design — materialize a Sendable snapshot with `array.materialized()` (or wrap a model in `MaterializedModule`) when you must cross a boundary.
 - **Don't use deprecated module imports**: Use `import MLX` not `import MLXRandom`.
 - **Don't forget to eval()**: Unevaluated arrays can accumulate large compute graphs.
 - **Don't mutate arrays directly**: Use operations that return new arrays.
@@ -415,6 +417,7 @@ See [deprecated.md](references/deprecated.md) for the complete migration guide.
 MLX has specific concurrency behavior:
 
 - **MLXArray is NOT Sendable**: This is intentional. Arrays contain references to compute graphs.
+- **Materialize to cross boundaries**: `array.materialized()` / `materialize(_:)` returns a `Sendable` `MaterializedArray` snapshot; `MaterializedModule` does the same for a whole model.
 - **evalLock protects eval/stream creation**: The global lock serializes evaluation and stream operations.
 - **Lazy operations are NOT thread-safe**: Don't share arrays across tasks without proper synchronization.
 - **Use actors to encapsulate MLX state**: Create and use MLXArrays within the same actor.
diff --git a/skills/mlx-swift/references/concurrency.md b/skills/mlx-swift/references/concurrency.md
index 96e4720a6..1b416c824 100644
--- a/skills/mlx-swift/references/concurrency.md
+++ b/skills/mlx-swift/references/concurrency.md
@@ -6,7 +6,7 @@ MLX Swift has specific concurrency characteristics for thread safety on Apple Si
 
 ### MLXArray is NOT Sendable
 
-`MLXArray` is intentionally **not** `Sendable`. This is by design:
+An ordinary `MLXArray` is intentionally **not** `Sendable`. This is by design:
 
 ```swift
 // WRONG: Will not compile in strict concurrency mode
@@ -20,6 +20,9 @@ Task {
 
 **Why?** MLXArray contains references to compute graphs that may be shared and mutated. Sending arrays across task boundaries could cause data races.
 
+The exception is a ``MaterializedArray`` (see below) — a fully-evaluated,
+immutable snapshot that *is* `Sendable`.
+
 ### Safe Patterns
 
 ```swift
@@ -40,6 +43,95 @@ Task {
     let newArray = MLXArray(data)
     // Work with newArray
 }
+
+// Pattern 3: Materialize a Sendable snapshot (no round-trip through [Float])
+let array = MLXArray([1, 2, 3])
+let snapshot = array.materialized()   // MaterializedArray: evaluated + Sendable
+
+Task {
+    print(snapshot + 3)   // safe to use across the task boundary
+}
+```
+
+## MaterializedArray and MaterializedModule
+
+When you actually need to hand an array (or a whole model) across a task or
+actor boundary, use the materialize APIs instead of copying data out to
+`[Float]` and rebuilding.
+
+### MaterializedArray
+
+`MaterializedArray` is a subclass of `MLXArray` that is a fully-evaluated,
+immutable snapshot and is declared `@unchecked Sendable`:
+
+- Contents are evaluated at construction, so no graph work is pending.
+- Mutating APIs (`_updateInternal`, `+=`, `-=`, `*=`, `/=`) are marked
+  `@available(*, unavailable)` and trap if called.
+- It is still an `MLXArray`, so it can be used anywhere one is accepted.
+  Operations involving it still produce ordinary (lazy) `MLXArray` results —
+  only the snapshot itself is frozen.
+
+Construct one via `MLXArray.materialized()` or the free `materialize(_:)`
+functions (never directly):
+
+```swift
+let m1 = a.materialized()          // single array
+let m2 = materialize(a)            // equivalent free function
+let batch = materialize([w, b])    // [MLXArray] -> [MaterializedArray], one eval
+
+// safe to send across boundaries
+let t = Task { print(m1 + 3) }
+```
+
+Use the `[MLXArray]` overload when materializing several arrays at once (for
+example a model's parameters) so MLX can schedule the eval as a single batch.
+
+### MaterializedModule
+
+`MaterializedModule<LayerType>` wraps a `Module` whose parameters have all been
+materialized, so the whole module is safe to share. A normal `Module` is not
+`Sendable` because its parameters are lazy, mutable `MLXArray` values.
+
+```swift
+// construct the base module inline so no other reference exists
+let lm = MaterializedModule(Linear(10, 10))
+
+let t = Task {
+    let x = MLXRandom.normal([10, 10])
+    print(sum(lm(x)))    // UnaryLayer call forwarded to the wrapped module
+}
+```
+
+Key points:
+
+- The initializer takes `base` as `consuming` — **do not retain or use the
+  original reference after passing it in.** Because `Module` is a reference
+  type this cannot be enforced by the compiler, so as a backstop the base and
+  all descendants are *sealed*: any later mutation (`update(parameters:)`,
+  `update(modules:)`, `apply(...)`, `freeze`/`unfreeze`, `train`) traps with
+  `fatalError` rather than silently violating the `Sendable` invariant.
+- Only parameters that can be replaced (e.g. wrapped with `@ParameterInfo`) are
+  actually retyped to `MaterializedArray`; all others are still evaluated and
+  materialized, they just keep the `MLXArray` static type.
+- `parameters()` returns a `NestedDictionary<String, MaterializedArray>`,
+  `parameterNBytes` gives the physical size (bytes) of all parameters, and
+  `logicalParameterCount` gives the number of logical parameters (quantized
+  layers are counted as their unquantized equivalent) for display purposes.
+- Calling the wrapped module is added per layer shape via extensions.
+  `UnaryLayer` support ships in the package; other protocols (e.g. a
+  `LanguageModel` protocol in `mlx-swift-lm`) are wired up with an
+  `extension MaterializedModule: Proto where LayerType: Proto` that forwards to
+  `_base` (accessed via `@_spi(MaterializedModule) import MLXNN`).
+
+This is the recommended way to store a model in a `Sendable` type:
+
+```swift
+struct ModelContext: Sendable {
+    var model: any LanguageModel & Sendable
+    init(model: some LanguageModel) {
+        self.model = MaterializedModule(model)
+    }
+}
 ```
 
 ## evalLock
@@ -97,6 +189,8 @@ These types are marked `@unchecked Sendable` and are safe to share:
 
 | Type | Location | Notes |
 |------|----------|-------|
+| `MaterializedArray` | MaterializedArray.swift | Evaluated, immutable array snapshot |
+| `MaterializedModule` | MaterializedModule.swift | Wraps a Module with materialized params |
 | `Stream` | Stream.swift | Immutable after creation |
 | `Device` | Device.swift | Immutable after creation |
 | `MLXFastKernel` | MLXFastKernel.swift | Kernel definition is immutable |
@@ -288,14 +382,17 @@ Why this is better:
 
 - Create MLXArrays within the task/actor that uses them
 - Use `eval()` before extracting values to send across boundaries
+- Use `materialized()` / `materialize(_:)` to get a `Sendable` snapshot when you must share an array
+- Wrap models in `MaterializedModule` to share them across concurrency domains
 - Use actors to encapsulate MLX state
-- Pass Sendable types (`[Float]`, `Int`, etc.) between tasks
+- Pass Sendable types (`[Float]`, `Int`, `MaterializedArray`, etc.) between tasks
 - Use explicit stream parameters for device/stream control
 - Use `WiredMemoryTicket.withWiredLimit` for concurrent GPU workloads
 
 ### DON'T
 
-- Share MLXArray instances across tasks
+- Share ordinary (lazy) MLXArray instances across tasks — materialize them first
+- Retain the original reference after passing a module to `MaterializedModule` (it is sealed and will trap on mutation)
 - Assume lazy operations are thread-safe
 - Create arrays in one task and use in another
 - Ignore the evalLock when doing custom threading

From 43c90af8b1f9de8ecf63d91f921e6efbdeb966e4 Mon Sep 17 00:00:00 2001
From: Lukas Zima <76228598+Silenterc@users.noreply.github.com>
Date: Wed, 21 Oct 2026 02:19:38 +0200
Subject: [PATCH 7/7] Fix tripletLoss + add tests (#505)

---
 Source/MLXNN/Losses.swift      | 15 +++++++-----
 Tests/MLXTests/LossTests.swift | 43 ++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/Source/MLXNN/Losses.swift b/Source/MLXNN/Losses.swift
index 218ae2544..bf8d8542e 100644
--- a/Source/MLXNN/Losses.swift
+++ b/Source/MLXNN/Losses.swift
@@ -226,7 +226,7 @@ public func smoothL1Loss(
 ///   - positives: positive samples
 ///   - negatives: negative samples
 ///   - axis: distribution axis
-///   - p: norm dree for pairwise distance
+///   - p: norm degree for pairwise distance
 ///   - margin: margin for the triplet loss
 ///   - eps: small positive constant to prevent numerical instability
 ///   - reduction: reduction type
@@ -238,13 +238,16 @@ public func tripletLoss(
     anchors: MLXArray, positives: MLXArray, negatives: MLXArray, axis: Int = -1, p: Int = 2,
     margin: Float = 1.0, eps: Float = 1e-6, reduction: LossReduction = .none
 ) -> MLXArray {
-    let loss = maximum(
-        sqrt(pow(anchors - positives, p).sum(axis: axis) + eps)
-            - sqrt(pow(anchors - negatives, p).sum(axis: axis) + eps)
-            + margin,
-        0
+    let positiveDistance = pow(
+        pow(abs(anchors - positives), p).sum(axis: axis) + eps,
+        1.0 / Float(p)
+    )
+    let negativeDistance = pow(
+        pow(abs(anchors - negatives), p).sum(axis: axis) + eps,
+        1.0 / Float(p)
     )
 
+    let loss = maximum(positiveDistance - negativeDistance + margin, 0)
     return reduction.reduce(loss: loss)
 }
 
diff --git a/Tests/MLXTests/LossTests.swift b/Tests/MLXTests/LossTests.swift
index b0efaeea0..8aba886af 100644
--- a/Tests/MLXTests/LossTests.swift
+++ b/Tests/MLXTests/LossTests.swift
@@ -29,4 +29,47 @@ class LossTests: XCTestCase {
         let loss2 = crossEntropy(logits: logits, targets: probs, reduction: .none)
         XCTAssertTrue(all(isNaN(loss2)).item())
     }
+
+    func testTripletLoss() {
+        // This mainly tests that tripletLoss now works for arbitrary values of p
+
+        let anchors = MLXArray([1, 2, 3, 1, 2, 3], [2, 3])
+        let positives = MLXArray([4, 5, 6, 0, -1, 2], [2, 3])
+        let negatives = MLXArray([7, 8, 9, 3, 2, 3], [2, 3])
+
+        // Test with reduction 'none'
+        let lossesNone = tripletLoss(
+            anchors: anchors, positives: positives, negatives: negatives, reduction: .none)
+        let expectedNone = MLXArray(converting: [0.0, 2.31662])
+        XCTAssertTrue(allClose(lossesNone, expectedNone).item())
+
+        // Test with reduction 'mean'
+        let lossesMean = tripletLoss(
+            anchors: anchors, positives: positives, negatives: negatives, reduction: .mean)
+        let expectedMean = mean(expectedNone)
+        XCTAssertTrue(allClose(lossesMean, expectedMean).item())
+
+        // Test with reduction 'sum'
+        let lossesSum = tripletLoss(
+            anchors: anchors, positives: positives, negatives: negatives, reduction: .sum)
+        let expectedSum = sum(expectedNone)
+        XCTAssertTrue(allClose(lossesSum, expectedSum).item())
+
+        // Test with non-default 'p' norm degrees
+        let normAnchors = MLXArray(converting: [1.0, 2.0]).reshaped([1, 2])
+        let normPositives = MLXArray(converting: [4.0, 6.0]).reshaped([1, 2])
+        let normNegatives = MLXArray(converting: [2.0, 3.0]).reshaped([1, 2])
+
+        let lossesP1 = tripletLoss(
+            anchors: normAnchors, positives: normPositives, negatives: normNegatives,
+            p: 1, reduction: .none)
+        let expectedP1 = MLXArray(converting: [6.0])
+        XCTAssertTrue(allClose(lossesP1, expectedP1).item())
+
+        let lossesP3 = tripletLoss(
+            anchors: normAnchors, positives: normPositives, negatives: normNegatives,
+            p: 3, reduction: .none)
+        let expectedP3 = MLXArray(converting: [4.23802])
+        XCTAssertTrue(allClose(lossesP3, expectedP3).item())
+    }
 }