MaterializedArray is a Sendable MLXArray by davidkoski · Pull Request #418 · ml-explore/mlx-swift

davidkoski · 2026-06-01T20:39:06Z

Proposed changes

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

Looking for feedback on this. I will add documentation once I test it a bit. Check out the tests to see how it works.

Checklist

Put an x in the boxes that apply.

I have read the CONTRIBUTING document
I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
I have added tests that prove my fix is effective or that my feature works
I have updated the necessary documentation (if needed)

davidkoski · 2026-06-01T20:39:41Z

+import Foundation
+import Numerics
+
+public final class MaterializedArray: MLXArray, @unchecked Sendable {


A new subtype of MLXArray that is Sendable

davidkoski · 2026-06-01T20:40:03Z

+
+public final class MaterializedArray: MLXArray, @unchecked Sendable {
+
+    init(materialized ctx: consuming mlx_array) {


Create using materialize(x) or x.materialized()

davidkoski · 2026-06-01T20:42:10Z

 import Numerics

-public final class MLXArray {
+public class MLXArray: ExpressibleByArrayLiteral {


remove final -- it is still closed outside the package

I don't think we were seeing any particular benefit from the final

add ExpressibleByArrayLiteral to the main type so that subclasses can override

davidkoski · 2026-06-01T20:45:17Z

+    }
+}
+
+extension MaterializedModule where LayerType: UnaryLayer {


We can use this pattern e.g. for LanguageModel in mlx-swift-lm

Added docs to show how this is done -- the some / any dance took a while to work out.

davidkoski · 2026-06-01T20:46:09Z

+    func update(
+        parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [],
+        modulePath: [String] = [],
+        mutate: (Module, String, MLXArray, MLXArray) -> Void


Normally we update the backing point in an MLXArray. For materialize() I want to replace the MLXArray instance with a MaterializedArray.

davidkoski · 2026-06-01T20:47:05Z

+    let x = MLXArray(10) + 5
+    let m = materialize(x)
+    let t = Task {
+        print(m + 3)


Tada, pass MLXArray (materialized) between isolation contexts! The compiler would give an error if we did this with x.

davidkoski · 2026-06-01T20:47:25Z

+
+    let t = Task {
+        let i = MLXRandom.normal([10, 10])
+        let r = lm(i)


The same thing with a Module.

- 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

manojmahapatra · 2026-06-17T07:52:43Z

+open class MaterializedModule<LayerType: Module>: Module, @unchecked Sendable {
+
+    /// Usable by extensions to implement `callAsFunction()` -- anyone else, DO NOT USE.
+    public let _base: LayerType


Question: would it make sense to avoid exposing this as public or at least rename it to something like _unsafeBase? Since callers can mutate the wrapper module through this reference.

Good question! It has to be public because:

the class can't be subclassed as it is final/Sendable (the former being a requirement for the latter)

since it can't be subclassed it has to be extended and the extensions can only access public members

But renaming it _unsafeBase is a good idea!

I have a staged commit where it marks the held model as immutable -- that will give additional runtime protection.

The integration in mlx-swift-lm for Embedders was good but perhaps too simple. I need to do the same for LLM/VLM and see what pops up.

This is now exposed via _spi so callers can't see it but implementations can.

manojmahapatra · 2026-06-18T06:05:12Z

+import MLX
+import MLXNN
+import Testing
+


Suggestion: It would be nice to add a unit test for MaterializedModule.parameters(), so it verify the wrapper still exposes the materialized parameter tree.

- 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

ronaldmannak · 2026-07-07T00:01:04Z

@davidkoski Is there any ETA on merging this PR? I'm working on a MLX-Swift-lm PR where this would be extremely useful

davidkoski · 2026-07-07T00:03:15Z

I think I have tested out most of this in my mlx-swift-lm branch (not all of it is pushed), so I can focus on wrapping this up. I will try to get this part in this week.

ronaldmannak · 2026-07-07T00:03:59Z

Great! Thanks David

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

davidkoski · 2026-07-08T17:39:57Z

-
-// MARK: - Expressible by literals
-
-extension MLXArray: ExpressibleByArrayLiteral {


This gets moved inside the class definition -- required for subclasses.

davidkoski · 2026-07-08T17:43:22Z

+/// because `UnaryLayer` is not a `Module`.
+///
+/// The class is `open` so you could give it a try.
+open class MaterializedModule<LayerType: Module>: IndentedDescription, @unchecked Sendable {


Note: this has been changed and is no longer a Module itself. That keeps the API surface smaller and means there are less confusing things that it exposes (like children() producing mutable Module instances).

In mlx-swift-lm (ml-explore/mlx-swift-lm#335) this worked out with a little surgery on the protocols.

Callers can still get a view of the parameters (now recast as MaterializedArray) and there is a common theme of counting the bytes for various reasons, now restored.

davidkoski · 2026-07-08T18:01:58Z

-    public func update(parameters: ModuleParameters) -> Self {
+    public func update(parameters: ModuleParameters) {


These are source breaking changes, though I expect they were never used. This was matching the python implementation:

def update(self, parameters: dict, strict: bool = True) -> Module

I guess this would allow a caller to chain calls, maybe layer.update(...).notSure(). It was never used in mlx-swift-lm.

Technically this is no longer required -- it was a side effect of the original implementation of MaterializedModule being is-a Module. I think it is a worthwhile cleanup so I am leaving it in.

- 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