On-device ML research with MLX and Swift

The Swift programming language has a lot of potential to be used for machine learning research because it combines the ease of use and high-level syntax of a language like Python with the speed of a compiled language like C++.

MLX is an array framework for machine learning research on Apple silicon. MLX is intended for research and not for production deployment of models in apps.

→ swift.org/blog/mlx-swift/

MLX Swift is a new open-source package from Apple that brings the MLX array framework for machine learning on Apple silicon to Swift. It gives researchers a Swift API to build and run models on Apple hardware.

The release includes a full Swift API for MLX core, higher-level neural network and optimiser packages, examples for text generation with Mistral 7B and MNIST training, and a C API as a bridge to the C++ core. All code uses a permissive MIT license.

MLX helps with machine learning research by offering hardware acceleration on CPU or GPU and automatic differentiation for training models. Swift fits well here because it is fast and easy to use on Apple silicon. With MLX Swift, researchers can experiment across platforms and devices.

To get started, set up with Xcode or SwiftPM. Basic array operations run on the default device, often the GPU. For example:

import MLX
import MLXRandom

let r = MLXRandom.normal([2])
print(r)  // array([-0.125875, 0.264235], dtype=float32)

let a = MLXArray(0 ..< 6, [3, 2])
print(a)  // array([[0, 1], [2, 3], [4, 5]], dtype=int32)

Function transformations let you compute gradients. For example:

func fn(_ x: MLXArray) -> MLXArray {
	x.square()
}

let gradFn = grad(fn)

let x = MLXArray(1.5)
let dfdx = gradFn(x)
print(dfdx)  // 3