Skip to content

MaterializedArray is a Sendable MLXArray #418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
davidkoski wants to merge 6 commits into
base: main
Choose a base branch
from
Open

MaterializedArray is a Sendable MLXArray #418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6c48e53
MaterializedArray is a Sendable MLXArray
davidkoski Jun 8, 2026
18106c8
add documentation
davidkoski Jun 8, 2026
ba24c6d
fix build issue
davidkoski Jun 8, 2026
0ceeada
no cross module links
davidkoski Jun 8, 2026
500ef2b
base needs to be public so other packages can extend
davidkoski Jun 8, 2026
a3fde72
relax the throw
davidkoski Jun 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Package.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// swift-tools-version: 5.12
// swift-tools-version: 6.2

@davidkoski davidkoski Jun 1, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Time to turn on the stricter concurrency checks.

// The swift-tools-version declares the minimum version of Swift required to build this package.
// Copyright © 2024 Apple Inc.

Expand Down
3 changes: 1 addition & 2 deletions Source/MLX/MLXArray+Indexing.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -388,8 +388,7 @@ extension MLXArray {
var result = mlx_array_new()
mlx_scatter(
&result, self.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx)
mlx_array_set(&self.ctx, result)
mlx_array_free(result)
self._updateInternal(MLXArray(result))
return
} else {
self._updateInternal(update)
Expand Down
36 changes: 0 additions & 36 deletions Source/MLX/MLXArray+Init.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -593,39 +593,3 @@ extension MLXArray {
}

}

// MARK: - Expressible by literals

extension MLXArray: ExpressibleByArrayLiteral {

// Note: MLXArray does not implement ExpressibleByFloatLiteral etc. because
// we want to create arrays in the context of the other arrays. For example:
//
// let x = MLXArray(1.5, dtype: .float16)
// let r = x + 2.5
//
// We expect r to have a dtype of float16. See ``ScalarOrArray``.

/// Initializer allowing creation of 1d `MLXArray` from an array literal.
///
/// ```swift
/// let a: MLXArray = [1, 2, 3]
/// ```
///
/// This is convenient for methods that have `MLXArray` parameters:
///
/// ```swift
/// print(array.take([1, 2, 3], axis: 0))
/// ```
///
/// ### See Also
/// - <doc:initialization>
public convenience init(arrayLiteral elements: Int32...) {
let ctx = elements.withUnsafeBufferPointer { ptr in
let shape = [Int32(elements.count)]
return mlx_array_new_data(
ptr.baseAddress!, shape, Int32(shape.count), Int32.dtype.cmlxDtype)
}
self.init(ctx)
}
}
93 changes: 71 additions & 22 deletions Source/MLX/MLXArray.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import Cmlx
import Foundation
import Numerics

public final class MLXArray {
public class MLXArray: ExpressibleByArrayLiteral {

@davidkoski davidkoski Jun 1, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • 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


/// Internal pointer to the mlx-c wrapper on `mlx::core::array`, used with `Cmlx` interop.
public internal(set) var ctx: mlx_array
Expand All @@ -19,6 +19,37 @@ public final class MLXArray {
self.ctx = ctx
}

// Note: MLXArray does not implement ExpressibleByFloatLiteral etc. because
// we want to create arrays in the context of the other arrays. For example:
//
// let x = MLXArray(1.5, dtype: .float16)
// let r = x + 2.5
//
// We expect r to have a dtype of float16. See ``ScalarOrArray``.

/// Initializer allowing creation of 1d `MLXArray` from an array literal.
///
/// ```swift
/// let a: MLXArray = [1, 2, 3]
/// ```
///
/// This is convenient for methods that have `MLXArray` parameters:
///
/// ```swift
/// print(array.take([1, 2, 3], axis: 0))
/// ```
///
/// ### See Also
/// - <doc:initialization>
required public convenience init(arrayLiteral elements: Int32...) {
let ctx = elements.withUnsafeBufferPointer { ptr in
let shape = [Int32(elements.count)]
return mlx_array_new_data(
ptr.baseAddress!, shape, Int32(shape.count), Int32.dtype.cmlxDtype)
}
self.init(ctx)
}

/// return the equivalent of a `.none` MLXArray (for the C API).
///
/// Not called `.none` to avoid ambiguity with `Optional`. This can be used
Expand All @@ -37,7 +68,7 @@ public final class MLXArray {
}

/// Number of bytes per element
public var itemSize: Int { mlx_array_itemsize(ctx) }
final public var itemSize: Int { mlx_array_itemsize(ctx) }

/// Total number of elements in the array
///
Expand All @@ -46,7 +77,7 @@ public final class MLXArray {
/// print(array.size)
/// // 12
/// ```
public var size: Int { mlx_array_size(ctx) }
final public var size: Int { mlx_array_size(ctx) }

/// Number of elements in the 0th dimension.
///
Expand All @@ -62,10 +93,10 @@ public final class MLXArray {
/// ...
/// }
/// ```
public var count: Int { dim(0) }
final public var count: Int { dim(0) }

/// Number of bytes in the array.
public var nbytes: Int { mlx_array_nbytes(ctx) }
final public var nbytes: Int { mlx_array_nbytes(ctx) }

/// Number of dimensions in the array.
///
Expand All @@ -74,7 +105,7 @@ public final class MLXArray {
/// print(array.ndim)
/// // 2
/// ```
public var ndim: Int { mlx_array_ndim(ctx) }
final public var ndim: Int { mlx_array_ndim(ctx) }

/// Data type of the elements in the array.
///
Expand All @@ -83,7 +114,7 @@ public final class MLXArray {
/// print(array.dtype)
/// // .int64 (aka Int.dtype)
/// ```
public var dtype: DType { DType(mlx_array_dtype(ctx)) }
final public var dtype: DType { DType(mlx_array_dtype(ctx)) }

/// Dimensions of the array.
///
Expand All @@ -92,7 +123,7 @@ public final class MLXArray {
/// print(array.shape)
/// // [3, 4]
/// ```
public var shape: [Int] {
final public var shape: [Int] {
let ndim = mlx_array_ndim(ctx)
guard ndim > 0 else { return [] }
let cShape = mlx_array_shape(ctx)!
Expand All @@ -104,7 +135,7 @@ public final class MLXArray {
/// ```swift
/// let (w, h) = array.shape2
/// ```
public var shape2: (Int, Int) {
final public var shape2: (Int, Int) {
let ndim = mlx_array_ndim(ctx)
precondition(ndim == 2)
let cShape = mlx_array_shape(ctx)!
Expand All @@ -116,7 +147,7 @@ public final class MLXArray {
/// ```swift
/// let (w, h, c) = array.shape3
/// ```
public var shape3: (Int, Int, Int) {
final public var shape3: (Int, Int, Int) {
let ndim = mlx_array_ndim(ctx)
precondition(ndim == 3)
let cShape = mlx_array_shape(ctx)!
Expand All @@ -128,7 +159,7 @@ public final class MLXArray {
/// ```swift
/// let (b, w, h, c) = array.shape4
/// ```
public var shape4: (Int, Int, Int, Int) {
final public var shape4: (Int, Int, Int, Int) {
let ndim = mlx_array_ndim(ctx)
precondition(ndim == 4)
let cShape = mlx_array_shape(ctx)!
Expand All @@ -149,7 +180,7 @@ public final class MLXArray {
/// Strides of the array backing.
///
/// Note: this is only stable once the array is evaluated.
var internalStrides: [Int] {
final var internalStrides: [Int] {
let ndim = mlx_array_ndim(ctx)
guard ndim > 0 else { return [] }
let strides = mlx_array_strides(ctx)!
Expand All @@ -167,7 +198,7 @@ public final class MLXArray {
/// // 4.5
/// let value: Float = array[1].item()
/// ```
public func item<T: HasDType>() -> T {
final public func item<T: HasDType>() -> T {
item(T.self)
}

Expand Down Expand Up @@ -328,7 +359,7 @@ public final class MLXArray {
/// // 4.5
/// let value = array[1].item(Float.self)
/// ```
public func item<T: HasDType>(_ type: T.Type) -> T {
final public func item<T: HasDType>(_ type: T.Type) -> T {
precondition(self.size == 1)
eval()

Expand Down Expand Up @@ -466,7 +497,7 @@ public final class MLXArray {
/// print(array.dim(1))
/// // 4
/// ```
public func dim(_ dim: Int) -> Int {
final public func dim(_ dim: Int) -> Int {
Int(mlx_array_dim(ctx, MLX.resolve(axis: dim, ndim: mlx_array_ndim(ctx)).int32))
}

Expand All @@ -481,7 +512,7 @@ public final class MLXArray {
/// print(array.dim(index))
/// // 4
/// ```
func dim(_ dim: Int32) -> Int32 {
final func dim(_ dim: Int32) -> Int32 {
mlx_array_dim(ctx, MLX.resolve(axis: Int(dim), ndim: mlx_array_ndim(ctx)).int32)
}

Expand All @@ -492,7 +523,7 @@ public final class MLXArray {
///
/// ### See Also
/// - <doc:conversion>
public func asType(_ type: DType, stream: StreamOrDevice = .default) -> MLXArray {
final public func asType(_ type: DType, stream: StreamOrDevice = .default) -> MLXArray {
guard type != self.dtype else { return self }
var result = mlx_array_new()
mlx_astype(&result, ctx, type.cmlxDtype, stream.ctx)
Expand All @@ -506,8 +537,9 @@ public final class MLXArray {
///
/// ### See Also
/// - <doc:conversion>
public func asType(_ type: (some HasDType).Type, stream: StreamOrDevice = .default) -> MLXArray
{
final public func asType(
_ type: (some HasDType).Type, stream: StreamOrDevice = .default
) -> MLXArray {
asType(type.dtype, stream: stream)
}

Expand All @@ -524,7 +556,7 @@ public final class MLXArray {
/// - ``realPart(stream:)``
/// - ``imaginaryPart(stream:)``
/// - <doc:conversion>
public func asImaginary(stream: StreamOrDevice = .default) -> MLXArray {
final public func asImaginary(stream: StreamOrDevice = .default) -> MLXArray {
precondition(!dtype.isComplex)
let i = MLXArray(real: 0, imaginary: 1)
return self * i
Expand All @@ -534,7 +566,7 @@ public final class MLXArray {
///
/// ### See Also
/// - <doc:conversion>
public func realPart(stream: StreamOrDevice = .default) -> MLXArray {
final public func realPart(stream: StreamOrDevice = .default) -> MLXArray {
precondition(dtype.isComplex)
return asType(Float.self)
}
Expand All @@ -543,7 +575,7 @@ public final class MLXArray {
///
/// ### See Also
/// - <doc:conversion>
public func imaginaryPart(stream: StreamOrDevice = .default) -> MLXArray {
final public func imaginaryPart(stream: StreamOrDevice = .default) -> MLXArray {
precondition(dtype.isComplex)
let i = MLXArray(real: 0, imaginary: 1)
return (self / i).asType(.float32)
Expand All @@ -559,6 +591,23 @@ public final class MLXArray {
}
}

/// Force evaluation and return a ``MaterializedArray`` snapshot of `self`.
///
/// The returned array is fully evaluated, immutable, and `Sendable` — see
/// ``MaterializedArray`` for the guarantees and trade-offs. Use this when
/// you need to hand an array across task or actor boundaries, or when you
/// need a stable reference that cannot be mutated by other code holding
/// the original `MLXArray`.
///
/// This is a thin convenience over ``materialize(_:)->MaterializedArray``.
///
/// ### See Also
/// - ``MaterializedArray``
/// - ``materialize(_:)->MaterializedArray``
public func materialized() -> MaterializedArray {
MLX.materialize(self)
}

/// Replace the contents with a reference to a new array (INTERNAL).
///
/// Note: this is an implementation detail and only visible because of the need to call it from
Expand Down
Loading
Loading