Skip to content

Add nested Swift Array initializers (and reverse) for MLXArray#409

Closed
john-rocky wants to merge 1 commit into
ml-explore:mainfrom
john-rocky:feat/issue-161-nested-array-init
Closed

Add nested Swift Array initializers (and reverse) for MLXArray#409
john-rocky wants to merge 1 commit into
ml-explore:mainfrom
john-rocky:feat/issue-161-nested-array-init

Conversation

@john-rocky
Copy link
Copy Markdown

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

  • `Source/MLX/MLXArray+NestedArrayInit.swift` — initializers and reverse extraction (purely Swift, no Cmlx calls beyond the existing 1-D path).
  • `Tests/MLXTests/MLXArray+NestedArrayInitTests.swift` — 12 tests covering dtype inference, shape, round-trip on plain 2/3/4-D, round-trip after `reshaped`, and cross-type extraction (`asArray2(Float.self)` from an int32 array).

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

  • 5-D and higher: per the issue thread ("the types will be tricky" for arbitrary N), and because 5+D literals are uncommon in user code.
  • Macro-based variant: @robertmsale floated an expression macro (`#mlx([[...]])`). That is a heavier dependency choice and orthogonal to this PR; can be added on top if desired.
  • `throws` instead of `precondition`: I matched the precondition style already used by `MLXArray(: [Int], _ shape:)` and `asArray(:)`. Happy to convert to `throws` if you prefer that for the shape-mismatch / rank-mismatch paths.

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.
@john-rocky
Copy link
Copy Markdown
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.

@john-rocky john-rocky closed this May 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support MLXArray initialization from nested Swift Arrays

1 participant