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

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

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

…e.g. quantized

- turn off training mode - clarify documentation - handle MaterializedArray passed as inputs (state) to a compiled function - contents are constant by definition so can be discarded and captured normally

davidkoski · 2026-08-19T22:48:08Z

@thechriswebb out-of-band PR feedback is addressed, thanks!

thechriswebb · 2026-08-20T19:02:48Z

+///
+/// ## Calling the wrapped module
+///
+/// `MaterializedModule`is not a `Module` subclass and cannot conform


nit: missing space

thechriswebb · 2026-08-20T19:05:49Z

+/// class MaterializedUnaryLayer: MaterializedModule<UnaryLayer> { ... }
+/// ```
+///
+/// because `UnaryLayer` is not a `Module`.


Is this right?

I should clarify that. It is wrong as written -- what is true is that it isn't a concrete type and you can't specialize with a protocol.

thechriswebb · 2026-08-20T19:07:14Z

+/// The initializer takes its base module as `consuming`:
+///
+/// ```swift
+/// let lm = try MaterializedModule(Linear(10, 10))


Outdated snippet? Init no longer throws

yep, good catch

thechriswebb · 2026-08-20T19:17:57Z

+    open override var parameterCount: Int {
+        if biases != nil {
+            scales.size * groupSize * 2
+        } else {
+            scales.size * groupSize
+        }
+    }


I may be reading this wrong, but I think scales.size * groupSize is already the unpacked weight count, so the doubling makes the number about twice as large as it should be. A quantized 1024x1024 layer reports 2097152 parameters, when the parameters it actually holds add up to 1082368. That is also more than the plain Linear it replaces, which reports 1049600.

I'm unsure about the condition too: biases is non-nil for every .affine layer, so it does not seem to separate out the layer bias. Did you mean to key on that instead?

My thinking was: biases is optional, but if present it is the same size as the scales.

But, I think what you are asking is about what this is trying to represent. This isn't the absolute byte count, it is a representation of what the non-quantized Linear represents.

So I think you are correct that it means to count the bias property instead, not what is physically stored.

Note: this is currently used in some size estimates for memory pressure and printing some stats.

thechriswebb · 2026-08-20T19:23:23Z

+    open override var parameterCount: Int {
+        scales.size * groupSize
+    }


Same question here. This one returns scales.size * groupSize while QuantizedLinear returns twice that for the same weights in .affine mode, so whichever way you want it counted, I think these two want to agree.

I think this is correct as-is -- Embedding only has a weight and we are trying to represent that.

davidkoski · 2026-08-20T20:12:46Z

@thechriswebb thank you for the feedback -- addressed in latest commit!