Discussion: Sendable MLXArray

[davidkoski]

MLXArray in mlx-swift is a class type and is not Sendable. The backing C++ object is not thread safe. Mostly. An evaluated MLXArray would be Sendable. I want to discuss some approaches we might take.

In all cases the type people typically interact with should be called MLXArray — this will preserve backward compatibility. The sendable version of MLXArray shouldn’t be mutated. The only ways to mutate something are vis in place operators like += or via _update() which is used internally e.g. via set on subscript. The Sendable array should be usable everywhere (except the mutating cases) that MLXArray can be used.

Ideally this would fit into Module as well, but that seems to be a bit more difficult.

The use case would be things like mlx-swift-lm where MLXArray is prepared on one Task and passed to another. We can guarantee that the arrays are evaluated (I think) so this could provide a nice benefit. Today we use some Box types behind the scene with promises in the comments (and the fact that they are private) giving some protection.

Sendable Subtype

I think this is the most likely approach (name is just a placeholder):

class MLXArray { /* not sendable */ }

class EvaluatedMLXArray: MLXArray, Sendable { … }

The current eval() function could be extended like this:

func eval(_ v: MLXArray) -> EvaluatedMLXArray { … }

Currently it doesn’t return anything so it could marked with @discardableResult.

Pros:

• the type that people interact with stays the same
• the subtype can be used anywhere that MLXArray can be used

Cons:

• this breaks Liskov substitution principle — the subtype restricts the operations
• this cannot be expressed in the type, e.g. removing subscript set or changing the += to not be in-place
• but it could be done at runtime

It does provide real benefit where parameters can be expressed in terms of EvaluatedMLXArray where they need to cross isolation context boundaries.

Abstract parent

class MLXArray { … }

/// this is roughly the current MLXArray
class MutableMLXArray: MLXArray { … }

class EvaluatedMLXArray: MLXArray, Sendable { … }

Pros:

• we could keep Liskov because MutableMLXArray and EvaluatedMLXArray could have different methods e.g. for += and subscript

Cons:

• it wouldn’t work because MLXArray doesn’t have all the same properties (maybe this could be dealt with because it would have most)
• initializers on MLXArray can’t return a MutableMLXArray, which would break current semantics

I don't see how we could make this work, but I think it is cleaner.

Protocol

protocol MLXArray { … }

/// this is roughly the current MLXArray
class MutableMLXArray: MLXArray { … }

class EvaluatedMLXArray: MLXArray, Sendable { … }

Same as Abstract Parent — MLXArray(10) would not work.

Additionally there are some difficulties in existential types. You might need any or some on the types. The performance aspect is probably a non-issue since these represent large computations to run on the GPU.