MLX Swift Numerical Computing

import MLX
import MLXRandom

// Generate noisy quadratic 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([1.0, -3.0, 2.0] as [Float]) // [bias, linear, quadratic]
let noise = MLXRandom.normal([x.shape[0]])
let y = trueTheta[0] + trueTheta[1] * x + trueTheta[2] * x * x + noise * 0.1

// Define the model: f(x; theta) = theta[0] + theta[1]*x + theta[2]*x²
func predict(_ xVals: MLXArray, _ theta: MLXArray) -> MLXArray {
    theta[0] + theta[1] * xVals + theta[2] * xVals * xVals
}

// Mean squared error loss
func loss(_ theta: MLXArray) -> MLXArray {
    let predictions = predict(x, theta)
    let diff = predictions - y
    return (diff * diff).mean()
}

// Transform loss into a gradient function w.r.t. theta
// MLX derives the gradient automatically — no hand-written derivatives
let gradFn = grad(loss)

// Initialize parameters and run gradient descent
var theta = MLXArray([0.0, 0.0, 0.0] as [Float])
let learningRate: Float = 0.05

for step in 0..<200 {
    let grads = gradFn(theta)
    theta = theta - learningRate * grads
    // eval() flushes the lazy compute graph each step
    eval(theta)
    
    if step % 50 == 0 {
        let currentLoss = loss(theta)
        eval(currentLoss)
        print("Step \(step): loss = \(currentLoss.item(Float.self))")
    }
}

// Read fitted coefficients (forces final computation)
let fitted = theta.asArray(Float.self)
print("Fitted: bias=\(fitted[0]), linear=\(fitted[1]), quadratic=\(fitted[2])")
// Expected approx: bias≈1.0, linear≈-3.0, quadratic≈2.0