Skip to content

Commit e613342

Browse files
committed
fix #364 -- x[.ellipsis] returns self instead of a "copy"
1 parent a2f0d76 commit e613342

File tree

2 files changed

+17
-1
lines changed

2 files changed

+17
-1
lines changed

Source/MLX/MLXArray+Indexing.swift

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -532,7 +532,9 @@ func getItem(src: MLXArray, operation: MLXArrayIndexOperation, stream: StreamOrD
532532
{
533533
switch operation {
534534
case .ellipsis:
535-
return src
535+
// let y = x[.elllipsis] -- this should be a "copy", e.g. a full range
536+
// slice. getItemND() handles that, so fall back to it
537+
return getItemND(src: src, operations: [operation], stream: stream)
536538

537539
case .newAxis:
538540
return src.expandedDimensions(axis: 0, stream: stream)

Tests/MLXTests/MLXArray+IndexingTests.swift

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -662,4 +662,18 @@ class MLXArrayIndexingTests: XCTestCase {
662662
assertEqual(b, expected)
663663
}
664664

665+
public func testCopyEllipsis() {
666+
// https://github.com/ml-explore/mlx-swift/issues/364
667+
let x = arange(0, 12).reshaped([3, 4])
668+
669+
// this should produce an independent copy of x -- they share the same backing
670+
// but if y is mutated then x should have reference to the original values
671+
let y = x[.ellipsis]
672+
673+
y[.ellipsis, 1] *= 10
674+
675+
// x and y should be different
676+
XCTAssertFalse(x.allClose(y).all().item())
677+
}
678+
665679
}

0 commit comments

Comments
 (0)