rename update -> _updateInternal to discourage callers (#230)

Author: davidkoski

Source/MLX/MLXArray+Indexing.swift

@@ -61,12 +61,6 @@ extension MLXArray {
         }
     }
 
-    /// Replace the interior ctx (`mlx_array` pointer) with a new value by transferring ownership
-    @inline(__always)
-    func update(_ ctx: mlx_array) {
-        mlx_array_set(&self.ctx, ctx)
-    }
-
     /// allow addressing as a positive index or negative (from end) using given axis
     @inlinable
     func resolve(index: Int, axis: Int) -> MLXArray {
@@ -161,7 +155,7 @@ extension MLXArray {
                 .broadcast(to: broadcastShape.asInt32)
 
             let indices = [resolve(index: index, axis: axis)]
-            self.update(scattered(indices: indices, updates: expanded, axes: [axis.int32]))
+            self._updateInternal(scattered(indices: indices, updates: expanded, axes: [axis.int32]))
         }
     }
 
@@ -350,7 +344,7 @@ extension MLXArray {
             let update = newValue.broadcast(to: broadcastShape).reshaped(updateShape)
 
             let axes = arange(axis + 1)
-            self.update(scattered(indices: arrayIndices, updates: update, axes: axes))
+            self._updateInternal(scattered(indices: arrayIndices, updates: update, axes: axes))
         }
     }
 
@@ -380,7 +374,7 @@ extension MLXArray {
         }
         set {
             if let result = updateSlice(src: self, operations: operations, update: newValue) {
-                self.update(result)
+                self._updateInternal(result)
                 return
             }
 
@@ -393,11 +387,11 @@ extension MLXArray {
                 var result = mlx_array_new()
                 mlx_scatter(
                     &result, self.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx)
-                self.update(result)
+                mlx_array_set(&self.ctx, result)
                 mlx_array_free(result)
                 return
             } else {
-                self.update(update)
+                self._updateInternal(update)
                 return
             }
         }

Source/MLX/MLXArray+Ops.swift

@@ -63,7 +63,7 @@ extension MLXArray {
     /// - <doc:arithmetic>
     /// - ``add(_:_:stream:)``
     public static func += (lhs: inout MLXArray, rhs: MLXArray) {
-        lhs.update(lhs + rhs)
+        lhs._updateInternal(lhs + rhs)
     }
 
     /// Element-wise addition with a ``ScalarOrArray`` (scalar) argument.
@@ -133,7 +133,7 @@ extension MLXArray {
     /// - <doc:arithmetic>
     /// - ``subtract(_:_:stream:)``
     public static func -= (lhs: inout MLXArray, rhs: MLXArray) {
-        lhs.update(lhs - rhs)
+        lhs._updateInternal(lhs - rhs)
     }
 
     /// Element-wise subtraction with a ``ScalarOrArray`` (scalar) argument.
@@ -224,7 +224,7 @@ extension MLXArray {
     /// - ``matmul(_:stream:)``
     /// - ``matmul(_:_:stream:)``
     public static func *= (lhs: inout MLXArray, rhs: MLXArray) {
-        lhs.update(lhs * rhs)
+        lhs._updateInternal(lhs * rhs)
     }
 
     /// Element-wise multiplication with a ``ScalarOrArray`` (scalar) argument.
@@ -340,7 +340,7 @@ extension MLXArray {
     /// - ``divide(_:_:stream:)``
     /// - ``floorDivide(_:_:stream:)``
     public static func /= (lhs: inout MLXArray, rhs: MLXArray) {
-        lhs.update(lhs / rhs)
+        lhs._updateInternal(lhs / rhs)
     }
 
     /// Element-wise division with a ``ScalarOrArray`` (scalar) argument.

Source/MLX/MLXArray.swift

@@ -539,8 +539,11 @@ public final class MLXArray {
         mlx_array_eval(ctx)
     }
 
-    /// Replace the contents with a reference to a new array.
-    public func update(_ array: MLXArray) {
+    /// Replace the contents with a reference to a new array (INTERNAL).
+    ///
+    /// 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)
     }
 

Source/MLX/Protocols.swift

@@ -5,7 +5,7 @@ import Foundation
 /// An object that can provide a list of the ``MLXArray`` in its inner state.
 ///
 /// Note that the array itself is not a reference to the inner state, but the ``MLXArray`` instances
-/// can be ``MLXArray/update(_:)`` to mutate the inner state.  The exact working is an
+/// can be ``MLXArray/_updateInternal(_:)`` to mutate the inner state.  The exact working is an
 /// implemention detail for MLX and should not be depended on by outside callers.
 ///
 /// ### See Also

Source/MLX/Transforms+Compile.swift

@@ -62,7 +62,7 @@ final class CompiledFunction: @unchecked (Sendable) {
 
             // replace the inner state with the tracers
             for (s, tracer) in zip(stateInputs, tracers[argumentsCount...]) {
-                s.update(tracer)
+                s._updateInternal(tracer)
             }
 
             // call the function with the tracer arguments
@@ -74,7 +74,7 @@ final class CompiledFunction: @unchecked (Sendable) {
 
             // put the original values back in the state
             for (s, saved) in zip(stateInputs, savedStateInputs) {
-                s.update(saved)
+                s._updateInternal(saved)
             }
 
             // return the result of the function and the state
@@ -106,7 +106,7 @@ final class CompiledFunction: @unchecked (Sendable) {
         let stateOutput = outputs.flatMap { $0.innerState() }
 
         for (s, newValues) in zip(stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) {
-            s.update(newValues)
+            s._updateInternal(newValues)
         }
 
         let resultLength = resultsPlusStateOutput.count - stateOutput.count

Source/MLXNN/Module.swift

@@ -448,7 +448,7 @@ open class Module {
                     throw UpdateError.mismatchedSize(
                         key: key, expectedShape: p.shape, actualShape: newArray.shape)
                 }
-                p.update(newArray)
+                p._updateInternal(newArray)
 
             case (.value(.parameters(let p)), .none):
                 if Self.parameterIsValid(key) {

Source/MLXNN/Normalization.swift

@@ -336,8 +336,8 @@ open class BatchNorm: Module, UnaryLayer {
 
         if self.training, let runningMean, let runningVar {
             let mu = momentum
-            runningMean.update((1 - mu) * runningMean + mu * mean)
-            runningVar.update((1 - mu) * runningVar + mu * variance)
+            runningMean._updateInternal((1 - mu) * runningMean + mu * mean)
+            runningVar._updateInternal((1 - mu) * runningVar + mu * variance)
 
         } else if let runningMean, let runningVar {
             mean = runningMean