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:40:49Z

-
-import Cmlx
-import Foundation
-import Numerics


Ignore this -- I had to move them to be part of MLXArray (not an extension) so that I could override them

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:43:17Z

+        self.init(ctx)
+    }
+
+    // MARK: - Inits


This is the body of MLXArray+Init.swift -- it is moved here because a subclass cannot override an init (or maybe even a method) that is not in the main class (vs extensions).

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

-    public let weight: MLXArray
-    public let bias: MLXArray?
+    @ParameterInfo public var weight: MLXArray
+    @ParameterInfo public var bias: MLXArray?


By making these @ParameterInfo we can replace the MLXArray with a MaterializedArray -- see the tests.

Note: this isn't required but it seemed like it might be useful.

davidkoski · 2026-06-01T20:44:48Z

+
+import MLX
+
+open class MaterializedModule<LayerType: Module>: Module, @unchecked Sendable {


A container for a Module that is Sendable.

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

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


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

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.

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

@@ -1,4 +1,4 @@
-// swift-tools-version: 5.12
+// swift-tools-version: 6.2


Time to turn on the stricter concurrency checks.

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

- 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.

manojmahapatra · 2026-06-18T05:54:31Z

+@Test
+func testMaterializedLinear() async throws {
+    let l = Linear(10, 10)
+    let lm = try MaterializedModule(l)


Nit: MaterializedModule's init is not throwing, so this test can drop both try and throws.

Yes, it was at one point but not now.

remove try

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

davidkoski commented Jun 1, 2026

View reviewed changes

davidkoski force-pushed the materialized-array branch from f18e7c2 to 2e0965e Compare June 8, 2026 05:11

davidkoski added 2 commits June 7, 2026 22:18

MaterializedArray is a Sendable MLXArray

6c48e53

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

add documentation

18106c8

davidkoski force-pushed the materialized-array branch from 2e0965e to 18106c8 Compare June 8, 2026 05:34

davidkoski added 4 commits June 7, 2026 22:41

fix build issue

ba24c6d

no cross module links

0ceeada

base needs to be public so other packages can extend

500ef2b

relax the throw

a3fde72

davidkoski mentioned this pull request Jun 9, 2026

Use MaterializedArray for Sendable conformance ml-explore/mlx-swift-lm#335

Open

4 tasks

manojmahapatra reviewed Jun 17, 2026

View reviewed changes

manojmahapatra reviewed Jun 18, 2026

View reviewed changes


		public final class MaterializedArray: MLXArray, @unchecked Sendable {

		init(materialized ctx: consuming mlx_array) {


		import MLX

		open class MaterializedModule<LayerType: Module>: Module, @unchecked Sendable {

		@@ -1,4 +1,4 @@
		// swift-tools-version: 5.12
		// swift-tools-version: 6.2

Uh oh!

Conversation

davidkoski commented Jun 1, 2026

Proposed changes

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

davidkoski Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

davidkoski Jun 1, 2026 •

edited

Loading