Use MaterializedArray for Sendable conformance

- adopt changes from ml-explore/mlx-swift#418
- we don't need private box types -- the technique becomes general
- it also opens up some potential for synchronous evaluation

Author: davidkoski

Libraries/MLXEmbedders/EmbedderModelContainer.swift

@@ -3,6 +3,8 @@
 import Foundation
 import MLXLMCommon
 
+// TODO dkoski -- make the example code be synchronous?
+
 /// Container for embedder models that guarantees single threaded access.
 ///
 /// Wrap models used by e.g. the UI in a ModelContainer. Callers can access
@@ -27,52 +29,52 @@ import MLXLMCommon
 /// }
 /// ```
 public final class EmbedderModelContainer: Sendable {
-    private let context: SerialAccessContainer<EmbedderModelContext>
+    private let context: EmbedderModelContext
 
     public var configuration: ModelConfiguration {
-        get async {
-            await context.read { $0.configuration }
-        }
+        context.configuration
     }
 
     public var tokenizer: Tokenizer {
-        get async {
-            await context.read { $0.tokenizer }
-        }
+        context.tokenizer
     }
 
     public var poolingStrategy: Pooling.Strategy {
-        get async {
-            await context.read { $0.pooling.strategy }
-        }
+        context.pooling.strategy
     }
 
     public init(context: consuming EmbedderModelContext) {
-        self.context = .init(context)
+        self.context = context
     }
 
     /// Perform an action on the ``EmbedderModelContext``.
     /// Callers _must_ eval any `MLXArray` before returning as `MLXArray` is not `Sendable`.
     ///
-    /// - Note: The closure receives `EmbedderModelContext` which is not `Sendable`. This is intentional -
-    ///   the closure runs within the actor's isolation, ensuring thread-safe access to the model.
     /// - Note: The `sending` keyword indicates the return value is transferred (not shared) across
     ///   isolation boundaries, allowing non-Sendable types to be safely returned.
     public func perform<R: Sendable>(
         _ action: @Sendable (EmbedderModelContext) async throws -> sending R
     ) async rethrows -> sending R {
-        try await context.read {
-            try await action($0)
-        }
+        try await action(context)
+    }
+
+    /// Perform an action on the ``EmbedderModelContext``.
+    ///
+    /// This is the synchronous form of ``perform(_:)`` and has
+    /// fewer restrictions.
+    public func perform<R>(
+        _ action: @Sendable (EmbedderModelContext) throws -> R
+    ) rethrows -> R {
+        try action(context)
     }
 
     @available(*, deprecated, message: "use perform(_: (EmbedderModelContext) -> R) instead")
     public func perform<R: Sendable>(
         _ action: @Sendable (EmbeddingModel, Tokenizer, Pooling) async throws -> sending R
     ) async rethrows -> sending R {
-        try await context.read {
-            try await action($0.model, $0.tokenizer, $0.pooling)
-        }
+        try await action(
+            context.model, context.tokenizer, context.pooling
+        )
     }
 
     /// Perform an action on the ``EmbedderModelContext`` with additional (non `Sendable`) context values.
@@ -83,32 +85,32 @@ public final class EmbedderModelContainer: Sendable {
         _ action: @Sendable (EmbedderModelContext, V) async throws -> R
     ) async rethrows -> sending R {
         let values = SendableBox(values)
-        return try await context.read {
-            try await action($0, values.consume())
-        }
+        return try await action(context, values.consume())
     }
 
     /// Update the owned `EmbedderModelContext`.
     /// - Parameter action: update action
+    @available(
+        *, unavailable,
+        message: "mutate EmbedderModelContext before passing to EmbedderModelContainer"
+    )
     public func update(_ action: @Sendable (inout EmbedderModelContext) -> Void) async {
-        await context.update {
-            action(&$0)
-        }
+        fatalError("update not supported")
     }
 
     // MARK: - Thread-safe convenience methods
 
     /// The resolved local model directory for the loaded container.
     public var modelDirectory: URL {
-        get async throws {
-            try (await configuration).modelDirectory
+        get throws {
+            try configuration.modelDirectory
         }
     }
 
     /// The resolved local tokenizer directory for the loaded container.
     public var tokenizerDirectory: URL {
-        get async throws {
-            try (await configuration).tokenizerDirectory
+        get throws {
+            try configuration.tokenizerDirectory
         }
     }
 }

Libraries/MLXEmbedders/EmbeddingModel.swift

@@ -48,3 +48,31 @@ extension EmbeddingModel {
             attentionMask: attentionMask)
     }
 }
+
+extension MaterializedModule: EmbeddingModel, BaseLanguageModel where LayerType: EmbeddingModel {
+
+    public var vocabularySize: Int { _base.vocabularySize }
+    public var poolingStrategy: Pooling.Strategy? { _base.poolingStrategy }
+    public var maxPositionEmbeddings: Int? { _base.maxPositionEmbeddings }
+
+    public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
+        _base.sanitize(weights: weights)
+    }
+
+    public func sanitize(weights: [String: MLXArray], metadata: [String: String]) -> [String:
+        MLXArray]
+    {
+        _base.sanitize(weights: weights, metadata: metadata)
+    }
+
+    public func callAsFunction(
+        _ inputs: MLXArray,
+        positionIds: MLXArray? = nil,
+        tokenTypeIds: MLXArray? = nil,
+        attentionMask: MLXArray? = nil
+    ) -> EmbeddingModelOutput {
+        return _base(
+            inputs, positionIds: positionIds, tokenTypeIds: tokenTypeIds,
+            attentionMask: attentionMask)
+    }
+}

Libraries/MLXEmbedders/ModelFactory.swift

@@ -126,18 +126,18 @@ public class EmbedderRegistry: AbstractModelRegistry, @unchecked Sendable {
 ///
 /// This is created using a ``EmbedderModelFactory`` and often used
 /// inside a ``EmbedderModelContainer``.
-public struct EmbedderModelContext {
+public struct EmbedderModelContext: Sendable {
     public var configuration: ModelConfiguration
-    public var model: any EmbeddingModel
+    public var model: any EmbeddingModel & Sendable
     public var tokenizer: any Tokenizer
     public let pooling: Pooling
 
     public init(
-        configuration: ModelConfiguration, model: any EmbeddingModel,
+        configuration: ModelConfiguration, model: some EmbeddingModel,
         tokenizer: any Tokenizer, pooling: Pooling
     ) {
         self.configuration = configuration
-        self.model = model
+        self.model = MaterializedModule(model)
         self.tokenizer = tokenizer
         self.pooling = pooling
     }

Libraries/MLXEmbedders/Pooling.swift

@@ -63,7 +63,7 @@ func loadPooling(modelDirectory: URL, model: EmbeddingModel) -> Pooling {
 ///
 /// `Pooling` takes the sequence of hidden states from a transformer model and collapses them
 /// into a single vector using strategies like mean, max, or token selection.
-open class Pooling: Module {
+public struct Pooling: Sendable {
 
     /// Supported pooling strategies.
     public enum Strategy: Sendable {
@@ -82,10 +82,10 @@ open class Pooling: Module {
     }
 
     /// The active strategy used for pooling hidden states.
-    public private(set) var strategy: Strategy
+    public let strategy: Strategy
 
     /// Optional dimension to truncate the resulting embedding to.
-    public private(set) var dimension: Int?
+    public let dimension: Int?
 
     /// Initializes a `Pooling` module with a specific strategy.
     /// - Parameters: