Swift Mlx

MLXOptimizers (Adam, AdamW, SGD, etc.)
       ↓
MLXNN (Layers, Modules, Losses)
       ↓
MLX (Arrays, Ops, Transforms, FFT, Linalg, Random)
       ↓
Cmlx (C/C++ bindings, Metal GPU)

import MLX

// Create arrays
let a = MLXArray([1, 2, 3, 4])
let b = MLXArray(0 ..< 12, [3, 4])  // Shape [3, 4]
let c = MLXArray.zeros([2, 3])
let d = MLXArray.ones([4, 4], dtype: .float32)

// Random arrays (use MLXRandom namespace or free functions)
let uniform = MLXRandom.uniform(0.0 ..< 1.0, [3, 3])
let normal = MLXRandom.normal([100])

let array = MLXArray(0 ..< 12, [3, 4])
array.shape    // [3, 4]
array.ndim     // 2
array.size     // 12
array.dtype    // .int64
array.count    // 3 (first dimension)

let a = MLXArray([1.0, 2.0, 3.0])
let b = MLXArray([4.0, 5.0, 6.0])

// Arithmetic (lazy - not computed until eval)
let sum = a + b
let product = a * b
let matmul = a.matmul(b.T)

// Force evaluation
eval(sum, product)
// or
sum.eval()

import MLX
import MLXNN

class MLP: Module, UnaryLayer {
    @ModuleInfo var fc1: Linear
    @ModuleInfo var fc2: Linear

    init(inputDim: Int, hiddenDim: Int, outputDim: Int) {
        self.fc1 = Linear(inputDim, hiddenDim)
        self.fc2 = Linear(hiddenDim, outputDim)
        super.init()
    }

    func callAsFunction(_ x: MLXArray) -> MLXArray {
        var x = fc1(x)
        x = relu(x)
        return fc2(x)
    }
}

let model = MLP(inputDim: 784, hiddenDim: 256, outputDim: 10)
eval(model)  // Initialize parameters

import MLXOptimizers

let model = MLP(inputDim: 784, hiddenDim: 256, outputDim: 10)
let optimizer = Adam(learningRate: 0.001)

func loss(model: MLP, x: MLXArray, y: MLXArray) -> MLXArray {
    let logits = model(x)
    return crossEntropy(logits: logits, targets: y, reduction: .mean)
}

// Compute loss and gradients - valueAndGrad returns a function
let lossAndGrad = valueAndGrad(model: model, loss)
let (lossValue, grads) = lossAndGrad(model, x, y)

// Update model
optimizer.update(model: model, gradients: grads)
eval(model, optimizer)

// Zeros and ones
MLXArray.zeros([3, 4])
MLXArray.ones([2, 2], dtype: .float16)

// Ranges
arange(0, 10, 2)           // [0, 2, 4, 6, 8]
linspace(0.0, 1.0, 5)      // [0.0, 0.25, 0.5, 0.75, 1.0]

// Identity and diagonal
MLXArray.identity(3)
diagonal(array, offset: 0)

// Full
MLXArray.full([2, 3], values: 7.0)

let a = MLXArray(0 ..< 12, [3, 4])

// Single element
a[0, 1]

// Slicing
a[0...]           // All rows
a[..<2]           // First 2 rows
a[1..., 2...]     // From row 1, column 2 onwards

// Advanced indexing
a[.ellipsis, 0]       // First column of all dimensions
a[.newAxis, .ellipsis]  // Add dimension at front

let a = MLXArray(0 ..< 12, [3, 4])

a.reshaped([4, 3])
a.reshaped(-1, 6)     // Infer first dimension
a.T                    // Transpose
a.transposed(1, 0)     // Explicit transpose
a.squeezed()           // Remove size-1 dimensions
a.expandedDimensions(axis: 0)

// Linear layers
Linear(inputDim, outputDim, bias: true)
Bilinear(in1, in2, out)

// Convolutions
Conv1d(inputChannels, outputChannels, kernelSize: 3)
Conv2d(inputChannels, outputChannels, kernelSize: 3, stride: 1, padding: 1)

// Normalization
LayerNorm(dimensions)
RMSNorm(dimensions)
BatchNorm(featureCount)
GroupNorm(groupCount, dimensions)

// Attention
MultiHeadAttention(dimensions: 512, numHeads: 8)

// Recurrent
RNN(inputSize, hiddenSize)
LSTM(inputSize, hiddenSize)
GRU(inputSize, hiddenSize)

// Regularization
Dropout(p: 0.1)

class MyLayer: Module {
    @ModuleInfo var layer: Linear           // Tracked module
    @ModuleInfo(key: "w") var weights: Linear  // Custom key

    let constant: MLXArray  // NOT tracked (no wrapper)
}

crossEntropy(logits: logits, targets: targets, reduction: .mean)
binaryCrossEntropy(logits: logits, targets: targets)
l1Loss(predictions: predictions, targets: targets, reduction: .mean)
mseLoss(predictions: predictions, targets: targets, reduction: .mean)
smoothL1Loss(predictions: predictions, targets: targets, beta: 1.0)
klDivLoss(inputs: inputs, targets: targets, reduction: .mean)

// Simple gradient
let gradFn = grad { x in
    sum(x * x)
}
let g = gradFn(MLXArray([1.0, 2.0, 3.0]))

// Value and gradient together
let (value, gradient) = valueAndGrad { x in
    sum(x * x)
}(MLXArray([1.0, 2.0, 3.0]))

// Model gradients - valueAndGrad returns a function, call it to get results
let lossAndGradFn = valueAndGrad(model: model) { model in
    model(input)
}
let (loss, grads) = lossAndGradFn(model)

// Common optimizers
let sgd = SGD(learningRate: 0.01, momentum: 0.9)
let adam = Adam(learningRate: 0.001, betas: (0.9, 0.999))
let adamw = AdamW(learningRate: 0.001, weightDecay: 0.01)

// Training step
optimizer.update(model: model, gradients: grads)
eval(model, optimizer)

// Compile a pure array function for faster execution
let compiledOp = compile { (a: MLXArray, b: MLXArray) -> MLXArray in
    let x = a + b
    return sum(x * x)
}

// Use compiled version
let output = compiledOp(arrayA, arrayB)

// Note: compile() works best with pure MLXArray functions.
// For models, call model methods directly (they can use internal compilation).

import MLX

let policy = WiredSumPolicy()

// Reservation: participates in admission but does not keep the wired limit high while idle.
let weightsTicket = policy.ticket(size: weightsBytes, kind: .reservation)
_ = await weightsTicket.start()

// Active work: raises limit while inference runs.
let inferenceTicket = policy.ticket(size: kvCacheBytes, kind: .active)
try await inferenceTicket.withWiredLimit {
    // run model inference
}

_ = await weightsTicket.end()