swift-mlx

MLX Swift Framework

MLX Swift is Apple's high-performance machine learning framework designed specifically for Apple Silicon. It provides NumPy-like array operations with lazy evaluation, automatic differentiation, and unified CPU/GPU memory.

When to Use This Skill

• Array operations on Apple Silicon (MLXArray)
• Building neural networks (MLXNN)
• Training models with automatic differentiation
• Custom Metal kernels via MLXFast
• Performance optimization with JIT compilation

Architecture Overview

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

Key File Reference

Quick Start

Basic Array Creation

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])

Array Properties

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)

DType Numeric Limits (finfo)

Floating-point limits for a DType, analogous to numpy.finfo. finfo is optional — nil for non-floating-point dtypes, so integer dtypes are rejected safely instead of returning a bogus value:

if let info = DType.float16.finfo {   // FInfo? — nil for non-float dtypes
    info.max               // 65504.0   (largest finite value)
    info.min               // -65504.0  (== -max)
    info.eps               // ulp of 1.0   (float16: 2^-10)
    info.smallestNormal    //              (float16: 2^-14)
    info.smallestSubnormal //              (float16: 2^-24)
}

DType.int32.finfo          // nil

When the dtype is statically known to be floating point, use the non-optional greatestFiniteMagnitude (mirrors Float.greatestFiniteMagnitude; traps on a non-float dtype):

let bound = DType.float16.greatestFiniteMagnitude   // 65504.0 (Double)

Attention Mask Fill (maskFill)

MLXArray.maskFill(for:) builds the value for masked-out positions before a softmax — -finfo(dtype).max, constructed directly in dtype so masked scores vanish under softmax with no asType to forget:

scores = MLX.where(causalMask, scores, MLXArray.maskFill(for: scores.dtype))

Requires a floating-point dtype (traps otherwise) — attention scores are float even when the KV cache is quantized.

Basic Operations

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()

Building a Neural Network

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

Training Loop

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)

Primary Workflow: Array Operations

See arrays.md for detailed array creation and indexing.

Creation Functions

// 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)

Indexing

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

Shape Manipulation

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)

Secondary Workflow: Neural Networks

See neural-networks.md for complete layer reference.

Built-in Layers

// 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)

Module Property Wrappers

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

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

Loss Functions

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)

Tertiary Workflow: Training

See transforms.md for automatic differentiation details.

Gradient Computation

// 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)

Optimizers

See optimizers.md for all optimizers.

// 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)

Compilation for Performance

// 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).

Quaternary Workflow: Wired Memory Coordination

See wired-memory.md for full policy, hysteresis, and admission guidance.

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()

Best Practices

DO

• Use lazy evaluation: MLX arrays are computed lazily. Call eval() strategically to control memory and compute.
• Batch eval calls: eval(a, b, c) is more efficient than separate calls.
• Use @ModuleInfo for all module properties to enable quantization and updates.
• Use actors for concurrent code: Encapsulate MLX state within actors for thread safety.
• Use namespaced functions: MLXRandom.uniform(), FFT.fft(), Linalg.inv().
• Use ticket-based wired memory coordination: Prefer WiredMemoryTicket.withWiredLimit and WiredMemoryManager.shared.

DON'T

• Don't share MLXArrays across tasks: MLXArray is NOT Sendable by design.
• **Don't use deprecated mo

Content truncated.