MLX Swift is an open-source numerical computing framework that brings NumPy-style n-dimensional array operations to Apple platforms, running automatically on the GPU with lazy evaluation and built-in automatic differentiation via `grad`.
• Write math-shaped code — array operations map directly to linear algebra notation without scalar loops or manual bookkeeping
• Automatic GPU execution and lazy evaluation deliver significant performance gains (10x+ over CPU scalar code) with no extra effort
• Built-in function transformations like `grad` enable automatic differentiation for gradient descent and ML model training without hand-written derivatives
Fits a quadratic polynomial to noisy data points using gradient descent and MLX's automatic differentiation (`grad`), demonstrating the core MLX Swift workflow: array ops, lazy eval, and `grad`.
import MLX
import MLXRandom
import Foundation
// MARK: - Generate noisy parabola data (y = 2x² - 3x + 1 + noise)
let x = MLXArray(stride(from: -2.0, through: 2.0, by: 0.1).map { Float($0) })
let trueTheta = MLXArray([2.0, -3.0, 1.0] as [Float]) // a, b, c
let noise = MLXRandom.normal([x.shape[0]])
let y = trueTheta[0] * x * x + trueTheta[1] * x + trueTheta[2] + noise * 0.1
// MARK: - Model: f(x; θ) = θ₀x² + θ₁x + θ₂
func predict(_ xVals: MLXArray, _ theta: MLXArray) -> MLXArray {
theta[0] * xVals * xVals + theta[1] * xVals + theta[2]
}
// MARK: - Loss: mean squared error
func mse(_ theta: MLXArray) -> MLXArray {
let residuals = predict(x, theta) - y
return (residuals * residuals).mean()
}
// MARK: - Automatic differentiation: derive gradient function from loss
let gradFn = grad(mse) // MLX derives ∂loss/∂theta automatically
// MARK: - Gradient descent optimisation loop
var theta = MLXArray([0.0, 0.0, 0.0] as [Float]) // start from zero
let lr: Float = 0.05
for step in 0..<200 {
let g = gradFn(theta) // evaluate gradient at current theta
theta = theta - lr * g // parameter update (lazy — not yet run)
MLX.eval(theta) // flush compute graph each iteration
if step % 40 == 0 {
let lossVal: Float = mse(theta).item()
let a: Float = theta[0].item()
let b: Float = theta[1].item()
let c: Float = theta[2].item()
print(String(format: "step %3d | loss %.4f | θ=[%.2f, %.2f, %.2f]",
step, lossVal, a, b, c))
}
}
// Final fitted parameters (should be close to [2, -3, 1])
let finalA: Float = theta[0].item()
let finalB: Float = theta[1].item()
let finalC: Float = theta[2].item()
print(String(format: "Fitted: %.3fx² + %.3fx + %.3f", finalA, finalB, finalC))MLX uses lazy evaluation — nothing executes until you call `MLX.eval()` or read a scalar value. In tight loops you must call `eval()` each iteration to prevent the compute graph from growing unboundedly. The package is installed via SPM, not bundled in the SDK.
Requires Apple Silicon or A-series GPU for full GPU acceleration; runs on CPU as fallback. Open-source via Swift Package Manager (github.com/ml-explore/mlx-swift).
More iOS 27 APIs land every week.
Get notified when new capabilities are published — no noise, just signal.