Add nested Swift Array initializers (and reverse) for MLXArray#409
Closed
john-rocky wants to merge 1 commit into
Closed
Add nested Swift Array initializers (and reverse) for MLXArray#409john-rocky wants to merge 1 commit into
john-rocky wants to merge 1 commit into
Conversation
Closes ml-explore#161. Adds row-major nested Swift Array initializers for `MLXArray` (and the matching reverse extraction methods) for 2-, 3-, and 4-dimensional arrays. Higher ranks are not commonly written as Swift literals and are intentionally left out — that question was the main reason this issue had been waiting (per @davidkoski's note: "We can certainly do it for a fixed number of dimensions"). 4-D is the upper bound chosen to cover `NCHW` / `NHWC` image inputs without introducing a macro dependency. ## Initializers added ```swift // 2-D MLXArray([[1.0, 2.0], [3.0, 4.0]]) // .float32, shape [2, 2] MLXArray([[1, 2], [3, 4]]) // .int32, shape [2, 2] (matches MLXArray([Int])) MLXArray(int64: [[1, 2], [3, 4]]) // .int64 MLXArray(converting: [[1.5, 2.5], [3.5, 4.5]]) // .float32 from Double // 3-D (same four shapes as above) MLXArray([[[1.0], [2.0]], [[3.0], [4.0]]]) // .float32, shape [2, 2, 1] // 4-D (same four shapes as above), useful for NCHW / NHWC literals MLXArray([[ [[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]], ]]) // .float32, shape [1, 2, 2, 2] ``` Each initializer flattens the nested array and uses the existing `MLXArray(_: [T], _ shape:)` path, so the dtype rules and behavior match the 1-D initializers exactly (notably: `[Int]` → `.int32`, `int64:` produces `.int64`, `converting:` casts `Double` → `Float`). Uneven inner-array lengths trigger a `precondition` failure, matching the existing convention for shape mismatches elsewhere in the file. ## Reverse extraction added @JimWallace asked in the issue for the inverse direction ("MLXArray to nested Swift arrays would also be handy"). Adding three matching helpers on `MLXArray`: ```swift let m = MLXArray([[1.0, 2.0], [3.0, 4.0]]) m.asArray2(Float.self) // [[1.0, 2.0], [3.0, 4.0]] m.reshaped([2, 2, 1]).asArray3(...) // 3-D nested m.reshaped([1, 2, 2, 1]).asArray4(...) // 4-D nested ``` `asArray2` / `asArray3` / `asArray4` reuse the existing `asArray(_:) -> [T]` path (so type conversion follows the same `asType` semantics — `asArray2(Float.self)` on an int32 array casts to float32). Rank mismatches trigger a `precondition` failure, matching the existing 1-D `asArray` style. ## Files - `Source/MLX/MLXArray+NestedArrayInit.swift` — initializers and reverse extraction. - `Tests/MLXTests/MLXArray+NestedArrayInitTests.swift` — 12 tests covering dtype inference, shape, round-trip, computed-tensor extraction, and cross-type extraction. ## Verification ``` swift build # Build complete!, no new warnings ``` `swift test` for the broader package hits the pre-existing "Failed to load the default metallib" error (ml-explore#349 / ml-explore#342) before any of the new tests can be exercised on this machine. The added code is pure Swift Array <-> existing `MLXArray(_: [T], _ shape:)` plumbing, so the runtime path is identical to the 1-D initializers, but I would appreciate a maintainer running the new test class through Xcode to confirm.
Author
|
Apologies — I missed @VDurocher's #392 when surveying open PRs and opened a duplicate. Their `MLXNestedArray` protocol approach supports arbitrary depth via the type system, which is a cleaner design than the fixed 2-D/3-D/4-D overloads I wrote here. Closing in favor of #392. If any of the round-trip tests against `reshaped` here are useful for #392, please feel free to lift them — no attribution needed. Sorry for the noise. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #161.
Adds row-major nested Swift Array initializers for
MLXArray(and the matching reverse extraction methods) for 2-, 3-, and 4-dimensional arrays. Higher ranks are intentionally left out — that question was the main reason this issue had been waiting (cf. @davidkoski: "We can certainly do it for a fixed number of dimensions… the types will be tricky" for arbitrary N). 4-D covers `NCHW` / `NHWC` image literals without introducing a macro dependency, so it felt like a natural stopping point.API additions
```swift
// 2-D
MLXArray([[1.0, 2.0], [3.0, 4.0]]) // .float32, shape [2, 2]
MLXArray([[1, 2], [3, 4]]) // .int32, shape [2, 2] (matches MLXArray([Int]))
MLXArray(int64: [[1, 2], [3, 4]]) // .int64
MLXArray(converting: [[1.5, 2.5], [3.5, 4.5]]) // .float32 from Double
// 3-D (same four shapes)
MLXArray([[[1.0], [2.0]], [[3.0], [4.0]]]) // .float32, shape [2, 2, 1]
// 4-D (same four shapes), useful for NCHW / NHWC literals
MLXArray([[
[[1.0, 2.0], [3.0, 4.0]],
[[5.0, 6.0], [7.0, 8.0]],
]]) // .float32, shape [1, 2, 2, 2]
```
Each initializer flattens the nested array and forwards to the existing `MLXArray(_: [T], _ shape:)` path, so the dtype rules match the 1-D initializers exactly (notably: `[Int]` → `.int32`, `int64:` produces `.int64`, `converting:` casts `Double` → `Float`). Uneven inner-array lengths trigger a `precondition` failure, matching the existing convention for shape-mismatch elsewhere in `MLXArray+Init.swift`.
Reverse extraction
@JimWallace asked in the issue for the inverse direction ("MLXArray to nested Swift arrays would also be handy"). Adding three matching helpers on
MLXArray:```swift
let m = MLXArray([[1.0, 2.0], [3.0, 4.0]])
m.asArray2(Float.self) // [[1.0, 2.0], [3.0, 4.0]]
m.reshaped([2, 2, 1]).asArray3(Float.self) // 3-D nested
m.reshaped([1, 2, 2, 1]).asArray4(Float.self) // 4-D nested
```
`asArray2` / `asArray3` / `asArray4` reuse the existing `asArray(_:) -> [T]` path, so cross-dtype extraction follows the same `asType` semantics — `asArray2(Float.self)` on an int32 array implicitly converts to float32. Rank mismatches trigger a `precondition` failure, matching the existing 1-D `asArray` style.
Files
Verification
```
swift build # Build complete!, no new warnings
```
`swift test` for the broader package hits the pre-existing "Failed to load the default metallib" error (#349 / #342) before any of the new tests can be exercised on this machine. The added code is pure Swift Array ↔ existing `MLXArray(_: [T], _ shape:)` plumbing, so the runtime path is identical to the 1-D initializers. I'd appreciate a maintainer running the new test class through Xcode to confirm; happy to iterate on the test set if helpful.
Intentional non-goals