diff --git a/Source/MLX/MLXArray+NestedArrayInit.swift b/Source/MLX/MLXArray+NestedArrayInit.swift new file mode 100644 index 000000000..15cc7b658 --- /dev/null +++ b/Source/MLX/MLXArray+NestedArrayInit.swift @@ -0,0 +1,331 @@ +// Copyright © 2026 Apple Inc. + +import Foundation + +// MARK: - Initialization from nested Swift Arrays + +extension MLXArray { + + // MARK: 2-D + + /// Create a 2-D ``MLXArray`` from a row-major ``[[T]]`` literal. + /// + /// Each inner array becomes a row of the resulting array. All rows must + /// have the same length. + /// + /// Example: + /// + /// ```swift + /// let a = MLXArray([[1.0, 2.0], + /// [3.0, 4.0]]) + /// // a.shape == [2, 2], a.dtype == .float32 + /// ``` + /// + /// Closes [#161](https://github.com/ml-explore/mlx-swift/issues/161). + /// + /// ### See Also + /// - + /// - ``asArray2(_:)`` + public convenience init(_ value: [[T]]) { + let rows = value.count + let cols = value.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == cols }, + "MLXArray(nested): rows must have equal length (got \(value.map(\.count)))" + ) + let flat = value.flatMap { $0 } + self.init(flat, [rows, cols]) + } + + /// Create a 2-D ``MLXArray`` of dtype `.int32` from a row-major ``[[Int]]`` literal. + /// + /// Mirrors the 1-D behavior of ``init(_:_:)-(value:[Int],_)``, which produces + /// an int32 array (and not int64) for ergonomic interop with Python / + /// numpy. + /// + /// ### See Also + /// - ``init(int64:)-([[Int]])`` + public convenience init(_ value: [[Int]]) { + let rows = value.count + let cols = value.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == cols }, + "MLXArray(nested): rows must have equal length (got \(value.map(\.count)))" + ) + let flat = value.flatMap { $0 } + self.init(flat, [rows, cols]) + } + + /// Create a 2-D `.int64` ``MLXArray`` from a row-major ``[[Int]]`` literal. + public convenience init(int64 value: [[Int]]) { + let rows = value.count + let cols = value.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == cols }, + "MLXArray(nested): rows must have equal length (got \(value.map(\.count)))" + ) + let flat = value.flatMap { $0 } + self.init(int64: flat, [rows, cols]) + } + + /// Create a 2-D `.float32` ``MLXArray`` from a row-major ``[[Double]]`` literal. + /// + /// `Double` is converted to `Float` because MLX does not support `.float64` + /// natively on all backends. + public convenience init(converting value: [[Double]]) { + let rows = value.count + let cols = value.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == cols }, + "MLXArray(nested): rows must have equal length (got \(value.map(\.count)))" + ) + let flat = value.flatMap { $0 } + self.init(converting: flat, [rows, cols]) + } + + // MARK: 3-D + + /// Create a 3-D ``MLXArray`` from a nested ``[[[T]]]`` literal in + /// outer-to-inner order. + /// + /// All sub-arrays at the same depth must have equal length. + /// + /// Example: + /// + /// ```swift + /// let a = MLXArray([[[1.0], [2.0]], + /// [[3.0], [4.0]]]) + /// // a.shape == [2, 2, 1] + /// ``` + public convenience init(_ value: [[[T]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == d1 && $0.allSatisfy { $0.count == d2 } }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0 } } + self.init(flat, [d0, d1, d2]) + } + + /// Create a 3-D `.int32` ``MLXArray`` from a ``[[[Int]]]`` literal. + public convenience init(_ value: [[[Int]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == d1 && $0.allSatisfy { $0.count == d2 } }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0 } } + self.init(flat, [d0, d1, d2]) + } + + /// Create a 3-D `.int64` ``MLXArray`` from a ``[[[Int]]]`` literal. + public convenience init(int64 value: [[[Int]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == d1 && $0.allSatisfy { $0.count == d2 } }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0 } } + self.init(int64: flat, [d0, d1, d2]) + } + + /// Create a 3-D `.float32` ``MLXArray`` from a ``[[[Double]]]`` literal. + public convenience init(converting value: [[[Double]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + precondition( + value.allSatisfy { $0.count == d1 && $0.allSatisfy { $0.count == d2 } }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0 } } + self.init(converting: flat, [d0, d1, d2]) + } + + // MARK: 4-D + + /// Create a 4-D ``MLXArray`` from a nested ``[[[[T]]]]`` literal. + /// + /// All sub-arrays at the same depth must have equal length. + /// + /// Useful for inputs shaped like `[batch, channels, height, width]` or + /// `[batch, height, width, channels]`. + public convenience init(_ value: [[[[T]]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + let d3 = value.first?.first?.first?.count ?? 0 + precondition( + value.allSatisfy { + $0.count == d1 + && $0.allSatisfy { + $0.count == d2 && $0.allSatisfy { $0.count == d3 } + } + }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0.flatMap { $0 } } } + self.init(flat, [d0, d1, d2, d3]) + } + + /// Create a 4-D `.int32` ``MLXArray`` from a ``[[[[Int]]]]`` literal. + public convenience init(_ value: [[[[Int]]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + let d3 = value.first?.first?.first?.count ?? 0 + precondition( + value.allSatisfy { + $0.count == d1 + && $0.allSatisfy { + $0.count == d2 && $0.allSatisfy { $0.count == d3 } + } + }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0.flatMap { $0 } } } + self.init(flat, [d0, d1, d2, d3]) + } + + /// Create a 4-D `.int64` ``MLXArray`` from a ``[[[[Int]]]]`` literal. + public convenience init(int64 value: [[[[Int]]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + let d3 = value.first?.first?.first?.count ?? 0 + precondition( + value.allSatisfy { + $0.count == d1 + && $0.allSatisfy { + $0.count == d2 && $0.allSatisfy { $0.count == d3 } + } + }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0.flatMap { $0 } } } + self.init(int64: flat, [d0, d1, d2, d3]) + } + + /// Create a 4-D `.float32` ``MLXArray`` from a ``[[[[Double]]]]`` literal. + public convenience init(converting value: [[[[Double]]]]) { + let d0 = value.count + let d1 = value.first?.count ?? 0 + let d2 = value.first?.first?.count ?? 0 + let d3 = value.first?.first?.first?.count ?? 0 + precondition( + value.allSatisfy { + $0.count == d1 + && $0.allSatisfy { + $0.count == d2 && $0.allSatisfy { $0.count == d3 } + } + }, + "MLXArray(nested): inner arrays must have equal length at every depth" + ) + let flat = value.flatMap { $0.flatMap { $0.flatMap { $0 } } } + self.init(converting: flat, [d0, d1, d2, d3]) + } +} + +// MARK: - Extracting nested Swift Arrays from an MLXArray + +extension MLXArray { + + /// Return the contents as a row-major 2-D Swift array. + /// + /// The array's ``dtype`` is preserved if it equals `T.dtype`; otherwise + /// values are cast to `T` (matching the semantics of ``asArray(_:)``). + /// The receiver must be a 2-D ``MLXArray``. + /// + /// Example: + /// + /// ```swift + /// let a = MLXArray([[1.0, 2.0], [3.0, 4.0]]) + /// let rows = a.asArray2(Float.self) + /// // rows == [[1.0, 2.0], [3.0, 4.0]] + /// ``` + /// + /// ### See Also + /// - + /// - ``asArray(_:)`` + /// - ``init(_:)-([[T]])`` + public func asArray2(_ type: T.Type) -> [[T]] { + precondition( + ndim == 2, + "asArray2 requires a 2-D MLXArray (got \(ndim)-D, shape \(shape))" + ) + let rows = dim(0) + let cols = dim(1) + let flat = asArray(type) + var result = [[T]]() + result.reserveCapacity(rows) + for r in 0..(_ type: T.Type) -> [[[T]]] { + precondition( + ndim == 3, + "asArray3 requires a 3-D MLXArray (got \(ndim)-D, shape \(shape))" + ) + let d0 = dim(0) + let d1 = dim(1) + let d2 = dim(2) + let flat = asArray(type) + var result = [[[T]]]() + result.reserveCapacity(d0) + for i in 0..(_ type: T.Type) -> [[[[T]]]] { + precondition( + ndim == 4, + "asArray4 requires a 4-D MLXArray (got \(ndim)-D, shape \(shape))" + ) + let d0 = dim(0) + let d1 = dim(1) + let d2 = dim(2) + let d3 = dim(3) + let flat = asArray(type) + var result = [[[[T]]]]() + result.reserveCapacity(d0) + for i in 0..