search

ml-explore / mlx-swift

Swift API for MLX
https://ml-explore.github.io/mlx-swift/
MIT License
427 stars 34 forks source link

Confusion about calling `asArray` on `MLXArray` returned from indexing #83

Open minghuaw opened 4 weeks ago

minghuaw commented 4 weeks ago

This may just be my confusion but as shown in the example below, I was expecting that calling asArray on the MLXArray returned from indexing would give [1, 2, 4, 5] but ended up getting [1,2,3,4]. This seems to be inconsistent with the python api where calling tolist() would give [[1, 2], [4, 5]].

The documentation doesn't seem to mention this behavior either

Return the contents as a single contiguous 1d Swift.Array.

let a = MLXArray(0..<9, [3,3])

let s = a[0..<2, 1..<3]
XCTAssertEqual(s.ndim, 2)
XCTAssertEqual(s.shape, [2,2])
assertEqual(s, MLXArray([1, 2, 4, 5], [2, 2]))

let s_arr = s.asArray(Int32.self)
XCTAssertEqual(s_arr, [1, 2, 3, 4])
davidkoski commented 4 weeks ago

I will take a look -- I suspect the readout (asArray) is expecting contiguous memory but not getting it. asData may have the same issue.

dcvz commented 4 weeks ago

I will take a look -- I suspect the readout (asArray) is expecting contiguous memory but not getting it. asData may have the same issue.

tolist() in python checks the stride: https://github.com/ml-explore/mlx/blob/21623156a32910b9db7c913f91612dcde0282caf/python/src/array.cpp#L34