feat: MLXArray initializers for nested Swift arrays (#161)

[VDurocher]

Closes #161

What

Adds MLXArray initializers that accept nested Swift arrays of any depth, matching the ergonomics of the Python mx.array([[1, 2], [3, 4]]) API.

How

• Introduces MLXNestedArray protocol with mlxShape and mlxFlattenedValues() requirements
• Array conforms when its Element also conforms (recursive) — supports any depth
• Scalar types (Float32, Float64, Int32, Int64, UInt8, UInt16, UInt32, Bool, Float16) conform as leaf types
• Dedicated MLXArray.init(_ rows: [[T]]) (2D) and MLXArray.init(_ slices: [[[T]]]) (3D) overloads for clean type inference at call-site
• Separate [[Int]] and [[[Int]]] overloads that mirror the existing [Int] behaviour (converts to Int32, with bounds-checking precondition)
• Shape validation via precondition — jagged arrays fail fast with a descriptive message

Examples

// 2D — shape: [2, 3], dtype: .float32
let matrix = MLXArray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

// 2D — shape: [2, 2], dtype: .int32  (mirrors existing [Int] → Int32 convention)
let ints = MLXArray([[7, 8], [9, 10]])

// 3D — shape: [2, 2, 2], dtype: .int32
let cube = MLXArray([[[1, 0], [0, 1]], [[2, 0], [0, 2]]])

// 3D — shape: [2, 3, 4], arbitrary element type
let tensor = MLXArray([[[Float16]]](…))

Files changed

File	Description
Source/MLX/MLXArray+NestedInit.swift	New — protocol + initialiseurs
Tests/MLXTests/MLXArray+NestedInitTests.swift	New — 16 test cases couvrant 2D, 3D, types scalaires, tableaux vides, indexation

Notes

• No macro magic — pure Swift generics
• No change to any existing initializer or public API
• The MLXNestedArray generic init is available as an escape hatch for deeper nesting (4D+), but the 2D/3D typed overloads are preferred for day-to-day use because Swift's type inference works better with concrete overloads
• [[Double]] works through the Float64: MLXNestedArray conformance (dtype .float64) — consistent with the existing init<T: HasDType>(_ value: [T]) behaviour for Double

[davidkoski]

Please translate comment to English

[davidkoski]

I like this approach!

[davidkoski]

I am not sure why I can't trigger CI to run, but this does fail swift-format -- please run:

pre-commit run --all-files

[davidkoski]

Do you need these overrides? The tests do not build for me, seeming because of them:

/Users/dkoski/Developer/mlx-swift/Tests/MLXTests/MLXArray+NestedInitTests.swift:19:22: error: ambiguous use of 'init(_:)'
        let matrix = MLXArray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] as [[Float32]])
                     ^
MLX.MLXArray.init:2:20: note: found this candidate in module 'MLX'
public convenience init<N>(_ nested: N) where N : MLX.MLXNestedArray, N : Collection, N.Element : MLX.MLXNestedArray}
                   ^
MLX.MLXArray.init:2:20: note: found this candidate in module 'MLX'
public convenience init<T>(_ rows: [[T]]) where T : MLX.HasDType}

If I comment all of these out, it does build but some tests fail.

[davidkoski]

Ideally we could have something like the MLXArray(converting:) that would handle this case.

[davidkoski]

Curiously this produces:

(lldb) p nested
([[Double]]) 2 values {
  [0] = 2 values {
    [0] = 7
    [1] = 8
  }
  [1] = 2 values {
    [0] = 9
    [1] = 10
  }
}

I think this is because the change is missing:

extension Int: MLXNestedArray {
    public typealias ScalarType = Int
    public var mlxShape: [Int] { [] }
    public func mlxFlattenedValues() -> [Int] { [self] }
}

so the next closest match is as a Double literal. (Int is not the same as Int64)