From f93decf4bcf23eed6bd3ca88bc296f555a9cfcaa Mon Sep 17 00:00:00 2001 From: ShankyShako <81709447+ShankyShako@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:42:44 -0500 Subject: [PATCH 1/3] Fix leak when _updateInternal replaces an array with siblings (#501) Fix leak when `_updateInternal` replaces an array with siblings `_updateInternal` assigned the new array with `mlx_array_set`, a C++ copy assignment that skips `array::~array()`. That destructor breaks the cycle between sibling outputs of a multi-output primitive. As a result, `Module.update(parameters:)` orphaned the init graph on every model load. This happened, for example, with a freshly built `QuantizedLinear` whose unevaluated random init weight is an output of `Quantize`. Now the code copies into a new handle and then frees the old one, so the destructor runs. Copying before freeing also keeps the contents intact when an array is passed to itself (`a._updateInternal(a)`), which the sharding path reaches. Co-authored-by: Claude Opus 5.5 --- Source/MLX/MLXArray.swift | 10 +++++++++- Tests/MLXTests/MLXArrayTests.swift | 7 +++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/Source/MLX/MLXArray.swift b/Source/MLX/MLXArray.swift index 8570cbac..bde1dd91 100644 --- a/Source/MLX/MLXArray.swift +++ b/Source/MLX/MLXArray.swift @@ -562,7 +562,15 @@ public final class MLXArray { /// Note: this is an implementation detail and only visible because of the need to call it from /// other `mlx-swift` modules. public func _updateInternal(_ array: MLXArray) { - mlx_array_set(&self.ctx, array.ctx) + // Free the old array rather than assigning over it: `mlx_array_set` is a C++ copy assignment, + // which skips `array::~array()` and with it the breaking of sibling cycles. An unevaluated + // output of a multi-output primitive (e.g. the `Quantize` behind a freshly built + // `QuantizedLinear`) replaced through `Module.update(parameters:)` would otherwise leak its + // whole graph. Copy before freeing so `a._updateInternal(a)` keeps its contents. + var new = mlx_array_new() + mlx_array_set(&new, array.ctx) + mlx_array_free(self.ctx) + self.ctx = new } /// Internal function for copying the backing `mlx::core::array` context. diff --git a/Tests/MLXTests/MLXArrayTests.swift b/Tests/MLXTests/MLXArrayTests.swift index 1a74fdba..2f556cbd 100644 --- a/Tests/MLXTests/MLXArrayTests.swift +++ b/Tests/MLXTests/MLXArrayTests.swift @@ -169,4 +169,11 @@ class MLXArrayTests: XCTestCase { } } + func testUpdateInternalWithItself() { + // `Module.update(parameters:)` can hand an array its own value back + let a = MLXArray([1, 2, 3] as [Float]) + a._updateInternal(a) + XCTAssertEqual(a.asArray(Float.self), [1, 2, 3]) + } + } From 31c03d9e3020eaa10e876df4e33c2ae38ab15874 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Fri, 16 Oct 2026 14:08:18 -0700 Subject: [PATCH 2/3] Enable JACCL (#502) --- Package.swift | 5 ++++- .../Cmlx/mlx-conditional/jaccl_conditional.cpp | 18 ++++++++++++++++++ Tests/MLXTests/DistributedTests.swift | 4 +++- xcode/MLX.xcodeproj/project.pbxproj | 1 + xcode/xcconfig/Cmlx.xcconfig | 2 +- 5 files changed, 27 insertions(+), 3 deletions(-) create mode 100644 Source/Cmlx/mlx-conditional/jaccl_conditional.cpp diff --git a/Package.swift b/Package.swift index ddc839fb..6188857d 100644 --- a/Package.swift +++ b/Package.swift @@ -209,6 +209,7 @@ let noCudaCmlxExcludes = [ // the wrappers there so they must not also be built directly "mlx/mlx/backend/cpu/compiled.cpp", "mlx/mlx/backend/cpu/jit_compiler.cpp", + "mlx/mlx/distributed/jaccl/no_jaccl.cpp", // opt-out of these backends (using metal) "mlx/mlx/backend/no_gpu", @@ -222,6 +223,7 @@ let noCudaCmlxExcludes = [ let cxxSettings: [CXXSetting] = [ .headerSearchPath("metal-cpp"), + .headerSearchPath("mlx/mlx/distributed/jaccl/lib"), .define("MLX_USE_ACCELERATE"), .define("ACCELERATE_NEW_LAPACK"), @@ -292,7 +294,8 @@ let cmlx = Target.target( "mlx/mlx/backend/metal/kernels", "mlx/mlx/backend/metal/nojit_kernels.cpp", - // distributed backends: enable ring, disable MPI + NCCL + JACCL + // distributed backends: enable ring, disable MPI + NCCL. JACCL is + // built through mlx-conditional on macOS, so its sources are excluded here. "mlx/mlx/distributed/mpi/mpi.cpp", "mlx/mlx/distributed/ring/no_ring.cpp", "mlx/mlx/distributed/nccl/nccl.cpp", diff --git a/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp b/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp new file mode 100644 index 00000000..8e8afea5 --- /dev/null +++ b/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp @@ -0,0 +1,18 @@ +// Copyright © 2026 Apple Inc. + +#include + +// JACCL (RDMA over Thunderbolt 5) needs , which ships with +// the macOS 26.2 SDK and later. Build the backend and its library when the +// header is present, and the stub otherwise. librdma is loaded at runtime, so +// on an older macOS MLXDistributed.isAvailable(.jaccl) reports false. +#if TARGET_OS_OSX && __has_include() +#include "../mlx/mlx/distributed/jaccl/jaccl.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/rdma.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/ring.cpp" +#include "../mlx/mlx/distributed/jaccl/lib/jaccl/tcp.cpp" +#else +#include "../mlx/mlx/distributed/jaccl/no_jaccl.cpp" +#endif diff --git a/Tests/MLXTests/DistributedTests.swift b/Tests/MLXTests/DistributedTests.swift index f9defc30..4946ea6b 100644 --- a/Tests/MLXTests/DistributedTests.swift +++ b/Tests/MLXTests/DistributedTests.swift @@ -42,7 +42,9 @@ class DistributedTests: XCTestCase { /// Backends that are not compiled in report `false` rather than trapping. func testUnavailableBackendsReportFalse() { XCTAssertFalse(MLXDistributed.isAvailable(.mpi)) - XCTAssertFalse(MLXDistributed.isAvailable(.jaccl)) + #if !os(macOS) + XCTAssertFalse(MLXDistributed.isAvailable(.jaccl)) + #endif // nccl is CUDA-only and is never built on Apple platforms #if !os(Linux) diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 0bb55a5a..00b532f5 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -1095,6 +1095,7 @@ mlx/distributed/jaccl/lib/jaccl/threadpool.h, mlx/distributed/jaccl/lib/jaccl/types.h, mlx/distributed/jaccl/lib/README.md, + mlx/distributed/jaccl/no_jaccl.cpp, mlx/distributed/mpi/CMakeLists.txt, mlx/distributed/mpi/mpi_declarations.h, mlx/distributed/mpi/mpi.cpp, diff --git a/xcode/xcconfig/Cmlx.xcconfig b/xcode/xcconfig/Cmlx.xcconfig index c6b2b5f4..57d2b911 100644 --- a/xcode/xcconfig/Cmlx.xcconfig +++ b/xcode/xcconfig/Cmlx.xcconfig @@ -5,7 +5,7 @@ PRODUCT_BUNDLE_IDENTIFIER = com.apple.mlx.Cmlx DEFINES_MODULE = YES BUILD_LIBRARY_FOR_DISTRIBUTION = YES -HEADER_SEARCH_PATHS = $(SDKROOT)/usr/include/c++/v1 $(SDKROOT)/usr/include $(inherited) $(SRCROOT)/../Source/Cmlx/metal-cpp $(SRCROOT)/../Source/Cmlx/fmt/include $(SRCROOT)/../Source/Cmlx/json/single_include/nlohmann +HEADER_SEARCH_PATHS = $(SDKROOT)/usr/include/c++/v1 $(SDKROOT)/usr/include $(inherited) $(SRCROOT)/../Source/Cmlx/metal-cpp $(SRCROOT)/../Source/Cmlx/fmt/include $(SRCROOT)/../Source/Cmlx/json/single_include/nlohmann $(SRCROOT)/../Source/Cmlx/mlx/mlx/distributed/jaccl/lib OTHER_CFLAGS = -isysroot $(SDKROOT) From 35c10b4954effc581963a066d5bf1fac817833dd Mon Sep 17 00:00:00 2001 From: Ahmed Hassan Date: Sat, 17 Oct 2026 04:19:01 +0500 Subject: [PATCH 3/3] Fix asArray crash on zero-size arrays (#486) Fix asArray crash on zero-size arrays An array with a logical size of 0 (e.g. an empty slice like a[128...] on a shorter axis) has no backing data pointer, so asArray() crashed by dereferencing nil. physicalSize now returns 0 for empty arrays, so asArray, asData and asDataCopy all return empty results from one place. The strided copy and the two no-copy asData paths also skip the nil pointer, which a transposed empty array or noCopy access would otherwise unwrap. This matches Python's tolist() behaviour. Adds tests for the empty-slice, empty-constructed, and scalar cases. Fixes #480 Co-authored-by: Claude Opus 5.5 --- Source/MLX/MLXArray+Bytes.swift | 31 ++++++++++++++++------- Tests/MLXTests/MLXArrayTests.swift | 40 ++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 9 deletions(-) diff --git a/Source/MLX/MLXArray+Bytes.swift b/Source/MLX/MLXArray+Bytes.swift index b4d1d62d..666f8a36 100644 --- a/Source/MLX/MLXArray+Bytes.swift +++ b/Source/MLX/MLXArray+Bytes.swift @@ -31,6 +31,13 @@ extension MLXArray { /// Return the physical size of the backing (assuming it is evaluated) in elements. This should /// only be used when accessing the backing directly, e.g. via `mlx_array_data_uint8()` var physicalSize: Int { + // an empty array (e.g. an empty slice like `a[128...]`) has no backing + // data pointer, so there is nothing to read even if some strides are + // non-zero + if self.size == 0 { + return 0 + } + // nbytes is the logical size of the input, not the physical size return zip(self.shape, self.internalStrides) .map { Swift.abs($0.0 * $0.1) } @@ -43,7 +50,11 @@ extension MLXArray { let shape = self.shape let strides = self.internalStrides - if contiguousDimension == 0 { + if output.isEmpty { + // an empty array has nothing to copy (and no backing pointer) + return + + } else if contiguousDimension == 0 { // entire backing is contiguous from.copyBytes(to: output) @@ -218,10 +229,11 @@ extension MLXArray { if self.contiguousToDimension() == 0 { // the backing is contiguous, we can provide a wrapper // for the contents without a copy - let source = UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx))! - let data = Data( - bytesNoCopy: source, count: size * itemSize, - deallocator: .none) + // an empty array has no backing pointer + let data = + UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx)).map { + Data(bytesNoCopy: $0, count: size * itemSize, deallocator: .none) + } ?? Data() return MLXArrayData( data: data, @@ -234,10 +246,11 @@ extension MLXArray { } case .noCopy: - let source = UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx))! - let data = Data( - bytesNoCopy: source, count: nbytes, - deallocator: .none) + // an empty array has no backing pointer + let data = + UnsafeMutableRawPointer(mutating: mlx_array_data_uint8(self.ctx)).map { + Data(bytesNoCopy: $0, count: nbytes, deallocator: .none) + } ?? Data() let strides: [Int] if ndim == 0 { diff --git a/Tests/MLXTests/MLXArrayTests.swift b/Tests/MLXTests/MLXArrayTests.swift index 2f556cbd..2d0dc427 100644 --- a/Tests/MLXTests/MLXArrayTests.swift +++ b/Tests/MLXTests/MLXArrayTests.swift @@ -36,6 +36,46 @@ class MLXArrayTests: XCTestCase { XCTAssertEqual(b, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) } + func testAsArrayEmpty() { + // an empty slice has size 0 and no backing pointer; it should return + // an empty array rather than crashing (see issue #480) + let a = MLXArray.ones([1, 200]) + let b = a[128...] + XCTAssertEqual(b.shape, [0, 200]) + XCTAssertEqual(b.size, 0) + XCTAssertEqual(b.asArray(Int.self), []) + XCTAssertEqual(b.asArray(Float.self), []) + + // a directly-constructed empty array behaves the same + let empty = MLXArray([Int32](), [0]) + XCTAssertEqual(empty.asArray(Int32.self), []) + + // a non-contiguous empty array goes through the strided copy path + let t = b.T + XCTAssertEqual(t.shape, [200, 0]) + XCTAssertEqual(t.asArray(Float.self), []) + } + + func testAsDataEmpty() { + // asData shares physicalSize with asArray, so it also returns empty + // data for an empty array (see issue #480) + let b = MLXArray.ones([1, 200])[128...] + for access in [MLXArray.AccessMethod.copy, .noCopyIfContiguous, .noCopy] { + let data = b.asData(access: access) + XCTAssertEqual(data.data.count, 0) + XCTAssertEqual(data.shape, [0, 200]) + } + XCTAssertEqual(b.T.asData(access: .copy).data.count, 0) + } + + func testAsArrayScalar() { + // a scalar (0-dimensional) array yields a single-element array + let a = MLXArray(7) + XCTAssertEqual(a.shape, []) + XCTAssertEqual(a.size, 1) + XCTAssertEqual(a.asArray(Int32.self), [7]) + } + func testAsArrayNonContiguous1() { // skipping elements via slicing let a = MLXArray(0 ..< 9, [3, 3])