Sync upstream ml-explore/mlx-swift main (35c10b4)

Package.swift

@@ -141,6 +141,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",
@@ -154,6 +155,7 @@ let noCudaCmlxExcludes = [
 
     let cxxSettings: [CXXSetting] = [
         .headerSearchPath("metal-cpp"),
+        .headerSearchPath("mlx/mlx/distributed/jaccl/lib"),
 
         .define("MLX_USE_ACCELERATE"),
         .define("ACCELERATE_NEW_LAPACK"),
@@ -224,7 +226,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",

Source/Cmlx/mlx-conditional/jaccl_conditional.cpp

@@ -0,0 +1,18 @@
+// Copyright © 2026 Apple Inc.
+
+#include <TargetConditionals.h>
+
+// JACCL (RDMA over Thunderbolt 5) needs <infiniband/verbs.h>, 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(<infiniband/verbs.h>)
+#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

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 {

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.

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)

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])
@@ -169,4 +209,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])
+    }
+
 }

xcode/MLX.xcodeproj/project.pbxproj

@@ -1113,6 +1113,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,

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)