🔄 Auto-Sync: ml-explore/mlx-swift main (5f117a2, 2 commit(s))

Source/MLX/Cmlx+Util.swift

@@ -10,6 +10,8 @@ func new_mlx_vector_array(_ arrays: some Collection<MLXArray>) -> mlx_vector_arr
     }
 }
 
+/// Convert `mlx_vector_array` into a Swift array of `MLXArray`.
+/// - Parameter vector_array: input vector
 func mlx_vector_array_values(_ vector_array: mlx_vector_array) -> [MLXArray] {
     (0 ..< mlx_vector_array_size(vector_array))
         .map { index in
@@ -20,6 +22,27 @@ func mlx_vector_array_values(_ vector_array: mlx_vector_array) -> [MLXArray] {
         }
 }
 
+/// Convert `mlx_vector_array` into a Swift array of `MLXArray` and ensure
+/// that there are at least `count` items in the result.
+///
+/// This variant is useful for callers that want to assume there are at
+/// least `count` values in the result, even though an error condition might
+/// result in fewer.
+///
+/// Note: the extra values supplied after the actual length of the vector
+/// are `mlxNone` and are not typically usable.
+///
+/// - Parameters:
+///   - vector_array: input vector
+///   - count: minimum number of values in the result
+func mlx_vector_array_values(
+    _ vector_array: mlx_vector_array, count: Int
+) -> [MLXArray] {
+    var result = mlx_vector_array_values(vector_array)
+    while result.count < count { result.append(.mlxNone) }
+    return result
+}
+
 func mlx_map_array_values(_ mlx_map: mlx_map_string_to_array) -> [String: MLXArray] {
     var result = [String: MLXArray]()
 

Source/MLX/Linalg.swift

@@ -223,7 +223,7 @@ public enum MLXLinalg {
         mlx_linalg_svd(&vec, array.ctx, true, stream.ctx)
         defer { mlx_vector_array_free(vec) }
 
-        let arrays = mlx_vector_array_values(vec)
+        let arrays = mlx_vector_array_values(vec, count: 3)
         return (arrays[0], arrays[1], arrays[2])
     }
 
@@ -242,7 +242,7 @@ public enum MLXLinalg {
         mlx_linalg_svd(&vec, array.ctx, false, stream.ctx)
         defer { mlx_vector_array_free(vec) }
 
-        let arrays = mlx_vector_array_values(vec)
+        let arrays = mlx_vector_array_values(vec, count: 1)
         return arrays[0]
     }
 
@@ -369,7 +369,7 @@ public enum MLXLinalg {
         var vec = mlx_vector_array_new()
         mlx_linalg_lu(&vec, a.ctx, stream.ctx)
         defer { mlx_vector_array_free(vec) }
-        let arrays = mlx_vector_array_values(vec)
+        let arrays = mlx_vector_array_values(vec, count: 3)
         return (arrays[0], arrays[1], arrays[2])
     }
 

Source/MLX/MLXArray+Indexing.swift

@@ -3,6 +3,8 @@
 import Cmlx
 import Foundation
 
+private let log = MLXLogger(label: "MLXArray+Indexing")
+
 private func int32Range(_ count: Int) -> [Int32] {
     Array(0 ..< Int32(count))
 }
@@ -334,6 +336,13 @@ extension MLXArray {
             // conform the shapes to each other
             arrayIndices = MLX.broadcast(arrays: arrayIndices)
 
+            guard !arrayIndices.isEmpty else {
+                // if we get to this point it means that there is an error
+                // handler that will catch the error
+                MLX.log.error("broadcast failed in subscript set")
+                return
+            }
+
             // the shape of the broadcast value
             let broadcastShape =
                 arrayIndices[0].shape.asInt32 + int32Range(axis + 1, self.ndim).map { dim($0) }
@@ -1154,7 +1163,16 @@ func scatterArguments(
     }
 
     // Broadcast the update to the indices and slices
-    arrayIndices = broadcast(arrays: arrayIndices)
+    arrayIndices = broadcast(arrays: arrayIndices, stream: stream)
+
+    guard !arrayIndices.isEmpty else {
+        // if we get to this point it means that there is an error
+        // handler that will catch the error
+        MLX.log.error("broadcast failed in scatterArguments")
+        // empty indices + src as the update -> caller does src._updateInternal(src), a no-op
+        return ([], src, [])
+    }
+
     let updateShapeBroadcast =
         arrayIndices[0].shape + sliceShapes + src.shape.dropFirst(nonNewAxisOperationCount)
     update = broadcast(update, to: updateShapeBroadcast, stream: stream)

Source/MLX/MLXArray+Ops.swift

@@ -2767,7 +2767,7 @@ extension MLXArray {
         var vec = mlx_vector_array_new()
         mlx_split(&vec, ctx, 2, axis.int32, stream.ctx)
         defer { mlx_vector_array_free(vec) }
-        let pieces = mlx_vector_array_values(vec)
+        let pieces = mlx_vector_array_values(vec, count: 2)
         return (pieces[0], pieces[1])
     }
 

Source/MLX/MLXFast.swift

@@ -36,7 +36,7 @@ public enum MLXFast {
             (mask ?? .mlxNone).ctx,
             stream.ctx)
 
-        let tuple = mlx_vector_array_values(result)
+        let tuple = mlx_vector_array_values(result, count: 2)
         return (output: tuple[0], state: tuple[1])
     }
 

Source/MLX/Ops+Array.swift

@@ -1684,7 +1684,7 @@ public func split(_ array: MLXArray, axis: Int = 0, stream: StreamOrDevice = .de
     var vec = mlx_vector_array_new()
     mlx_split(&vec, array.ctx, 2, axis.int32, stream.ctx)
     defer { mlx_vector_array_free(vec) }
-    let pieces = mlx_vector_array_values(vec)
+    let pieces = mlx_vector_array_values(vec, count: 2)
     return (pieces[0], pieces[1])
 }
 

Source/MLX/Ops.swift

@@ -1066,7 +1066,7 @@ public func depends(input: MLXArray, dependencies: [MLXArray]) -> MLXArray {
     defer { mlx_vector_array_free(dependencies) }
     mlx_depends(&result, inputs, dependencies)
 
-    let arrays = mlx_vector_array_values(result)
+    let arrays = mlx_vector_array_values(result, count: 1)
     return arrays[0]
 }
 
@@ -1235,7 +1235,7 @@ public func divmod(
     var vec = mlx_vector_array_new()
     mlx_divmod(&vec, a.ctx, b.ctx, stream.ctx)
     defer { mlx_vector_array_free(vec) }
-    let result = mlx_vector_array_values(vec)
+    let result = mlx_vector_array_values(vec, count: 2)
     return (result[0], result[1])
 }
 
@@ -2510,7 +2510,7 @@ public func quantized(
         (globalScale ?? .mlxNone).ctx,
         stream.ctx)
 
-    let arrays = mlx_vector_array_values(r)
+    let arrays = mlx_vector_array_values(r, count: 2)
     return (arrays[0], arrays[1], arrays.count > 2 ? arrays[2] : nil)
 }
 
@@ -3061,7 +3061,7 @@ public func uniqued(
         (fillValue ?? .mlxNone).ctx,
         stream.ctx)
 
-    let arrays = mlx_vector_array_values(vector)
+    let arrays = mlx_vector_array_values(vector, count: 1)
     return arrays[0]
 }
 
@@ -3159,7 +3159,9 @@ public func uniqued(
         (fillValue ?? .mlxNone).ctx,
         stream.ctx)
 
-    let arrays = mlx_vector_array_values(vector)
+    let expectedCount =
+        1 + [UniqueOptions.index, .inverse, .counts].filter { returning.contains($0) }.count
+    let arrays = mlx_vector_array_values(vector, count: expectedCount)
     var result = UniqueResult(values: arrays[0])
 
     var readIndex = 1

Source/MLX/Transforms+Compile.swift

@@ -47,7 +47,7 @@ final class CompiledFunction: @unchecked (Sendable) {
 
     /// guards the observed state (``inputs`` / ``outputs``) -- see #226
     ///
-    /// This is always acquired *inside* `evalLock`; see ``call(_:)`` for the
+    /// This is always acquired *inside* `evalLock`; see ``call(_:expected:)`` for the
     /// ordering rule.  It is recursive because a compiled function's body can
     /// legitimately re-enter that same compiled function on the same thread
     /// while it is being traced (a nested call with tracer inputs bypasses the
@@ -156,8 +156,15 @@ final class CompiledFunction: @unchecked (Sendable) {
     /// hold it through `mlx_closure_apply` -- i.e. across the actual evaluation
     /// -- so this only moves the acquisition earlier by the state gathering and
     /// closure creation at the top of ``innerCall(_:)``.
-    func call(_ arguments: [MLXArray]) -> [MLXArray] {
-        withEvalLock {
+    ///
+    /// - Parameters:
+    ///   - arguments: arguments to the compiled function
+    ///   - expected: minimum number of results.  If an MLX error fires inside a
+    ///     `withError` scope the error is stored in the ErrorBox and the results
+    ///     are padded with `mlxNone` so callers can index them without crashing
+    ///     and `withError` can throw.
+    func call(_ arguments: [MLXArray], expected: Int) -> [MLXArray] {
+        var results = withEvalLock {
             // Stateless functions need only evalLock for cache bookkeeping.
             if inputs.isEmpty && outputs.isEmpty {
                 return innerCall(arguments)
@@ -166,11 +173,13 @@ final class CompiledFunction: @unchecked (Sendable) {
                 innerCall(arguments)
             }
         }
+        while results.count < expected { results.append(.mlxNone) }
+        return results
     }
 
     /// Acquire ``lock``, checking the ordering invariant first.
     ///
-    /// The check is where the invariant actually lives, not in ``call(_:)``: a
+    /// The check is where the invariant actually lives, not in ``call(_:expected:)``: a
     /// change that moved the per-instance lock back outside `evalLock` -- the
     /// shape that produced the deadlock -- trips this on the first compiled
     /// call, deterministically and without needing to lose a race.
@@ -184,7 +193,7 @@ final class CompiledFunction: @unchecked (Sendable) {
     }
 
     /// - Precondition: `evalLock` is held by the caller, and ``lock`` is held
-    ///   too unless this function has no observed state; see ``call(_:)``,
+    ///   too unless this function has no observed state; see ``call(_:expected:)``,
     ///   which is the only caller.
     func innerCall(_ arguments: [MLXArray]) -> [MLXArray] {
         let cache = ThreadCompilationCache.current
@@ -260,7 +269,7 @@ final class CompiledFunction: @unchecked (Sendable) {
         // note: this will use the cached compile (via the id)
         // but will be able to re-evaluate with fresh state if needed
         //
-        // evalLock is already held by call(_:) and covers everything through
+        // evalLock is already held by call(_:expected:) and covers everything through
         // mlx_closure_apply below, which is where the trace runs.
         var compiled = mlx_closure_new()
         let compileStatus = mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0)
@@ -308,6 +317,13 @@ final class CompiledFunction: @unchecked (Sendable) {
         // push the stateOutput into the state
         let stateOutput = outputs.flatMap { $0.innerState() }.filter { !$0.isSealed }
 
+        // a short result means something went wrong -- don't push result
+        // arrays into the state (and avoid a negative prefix below).
+        // call(_:expected:) pads the results.
+        guard resultsPlusStateOutput.count >= stateOutput.count else {
+            return []
+        }
+
         for (s, newValues) in zip(stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) {
             s._updateInternal(newValues)
         }
@@ -344,7 +360,7 @@ public func compile(
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless, f)
 
     return { arrays in
-        compileState.call(arrays)
+        compileState.call(arrays, expected: 0)
     }
 }
 
@@ -369,7 +385,7 @@ public func compile(
     let compileState = CompiledFunction(shapeless: shapeless, f)
 
     return { arrays in
-        compileState.call(arrays)
+        compileState.call(arrays, expected: 0)
     }
 }
 
@@ -388,11 +404,7 @@ public func compile(
     }
 
     return { a in
-        let r = compileState.call([a])
-        // r is empty only when an MLX error fired inside a withError scope — the
-        // error is already stored in the ErrorBox.  Return a placeholder so that
-        // withError can throw instead of crashing with "Index out of range".
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a], expected: 1)[0]
     }
 }
 
@@ -411,11 +423,7 @@ public func compile(
     }
 
     return { a in
-        let r = compileState.call([a])
-        // r is empty only when an MLX error fired inside a withError scope — the
-        // error is already stored in the ErrorBox.  Return a placeholder so that
-        // withError can throw instead of crashing with "Index out of range".
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a], expected: 1)[0]
     }
 }
 
@@ -436,8 +444,7 @@ public func compile(
     }
 
     return { a, b in
-        let r = compileState.call([a, b])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b], expected: 1)[0]
     }
 }
 
@@ -458,8 +465,7 @@ public func compile(
     }
 
     return { a, b in
-        let r = compileState.call([a, b])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b], expected: 1)[0]
     }
 }
 
@@ -480,8 +486,7 @@ public func compile(
     }
 
     return { a, b, c in
-        let r = compileState.call([a, b, c])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b, c], expected: 1)[0]
     }
 }
 
@@ -502,8 +507,7 @@ public func compile(
     }
 
     return { a, b, c in
-        let r = compileState.call([a, b, c])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b, c], expected: 1)[0]
     }
 }