WWDC 2026: Explore numerical computing in Swift with MLX (Session 328)

Use MLX Swift for NumPy-style n-dimensional arrays, lazy GPU execution, convolutions, and automatic differentiation in Swift numerical code.

Watch on Apple Developer

TL;DR

MLX Swift is positioned for mathematical Swift code that operates on n-dimensional arrays, with lazy evaluation enabling automatic GPU execution and function transformations such as grad.
The session compares scalar Swift with MLX Swift for Mandelbrot computation, showing how array operations express the math directly and run across the grid on the GPU.
Heat distribution is modeled with Jacobi iteration using conv2d, then accelerated with Successive Over-Relaxation using an omega parameter and red/black checkerboard masks.
A curve-fitting example uses automatic differentiation to compute gradients for a mean-squared-error loss and run a gradient descent optimization loop without hand-written derivatives.

MLX Swift is for numerical computing code where the primary goal is expressing mathematical operations clearly while still getting performance. Apple's broader numerical ecosystem still matters: Accelerate provides tuned CPU vector primitives, BNNS provides neural-network building blocks, Metal Performance Shaders provides direct GPU kernels, and Swift Numerics provides types such as Complex and generic numeric protocols.

MLX Swift's central abstraction is the n-dimensional array, similar to NumPy. The framework is open source under the MIT license, and the same MLX concepts are available across Swift, Python, C++, and C front ends.

Use MLX Swift when array-oriented mathematical code is more important than low-level control over kernels or memory bookkeeping.
Most NumPy-style code patterns translate to MLX Swift with minimal conceptual changes.
Lazy evaluation builds a compute graph and executes when a value is read or eval is called.
The lazy model also supports automatic GPU execution and automatic differentiation.

Core array model and lazy evaluation

The power-iteration example demonstrates how matrix-vector algorithms map directly into MLX Swift operations. MLXArray operations such as transpose, matrix multiplication, vector normalization, and addition look close to the mathematical notation.

Because operations are lazy, loops should call eval when needed to keep the graph from growing without bound. Reading a scalar result also forces computation.

.T gives a transpose.
matmul performs matrix-vector or matrix-matrix multiplication.
norm computes the vector norm used for normalization.
Use MLX's linear algebra package when full eigensystem routines are needed rather than implementing an iterative method manually.

Power iteration with MLX Swift arrays

Shows array operations, transposition, matrix multiplication, normalization, and explicit evaluation inside an iterative loop.

import MLX

let n = 100
let steps = 10
let B = MLXRandom.normal([n, n])
var v = MLXRandom.normal([n])

// Symmetric matrix A = Bᵀ + B
let A = B.T + B

// v ← A v / ‖A v‖
for _ in 0 ..< steps {
    let Av = matmul(A, v)
    v = Av / norm(Av)
    eval(v)
}

// λ = vᵀ A v
let lambda = matmul(matmul(v.T, A), v)
print(lambda)

Array computing with Mandelbrot

The Mandelbrot set is a useful example because the same recurrence, z = z² + c, is applied independently across a large grid of complex values. A scalar Swift implementation loops over pixels and iterations manually; the MLX Swift version constructs the whole grid and applies the recurrence to every point at once.

The MLX version removes per-pixel bookkeeping and uses the GPU by default. The exact speedup depends on the algorithm and workload, but the session notes that a 10x improvement is possible for this kind of array-parallel computation.

Plain Swift is expressive but scalar-at-a-time for this implementation.
MLX Swift expresses the grid as arrays and applies elementwise complex operations.
The escape count is accumulated with array comparisons rather than nested pixel management.

Mandelbrot set in MLX Swift

Builds a complex grid and applies the Mandelbrot recurrence across the full array instead of iterating scalar pixels manually.

import MLX

let x = linspace(-2.0, 0.5, count: w)
let y = linspace(-1.25, 1.25, count: h).reshaped(h, 1)
let c = x + y.asImaginary()

var z = MLXArray.zeros(like: c)
var counts = MLXArray.zeros(c.shape, dtype: .int16)

for _ in 0 ..< maxIterations {
    z = z * z + c
    counts = counts + (abs(z) .< 2)
}

Stencil computation with convolutions

The heat-distribution example models a room as a 2D temperature grid. Jacobi iteration updates each cell by averaging its four neighbors, which is exactly a stencil operation and can be implemented as a convolution.

Boundary conditions are handled with which, an elementwise ternary operation. Heat sources and walls keep fixed values; other cells use the next value computed by conv2d.

Represent the four-neighbor averaging stencil as a 3×3 convolution kernel.
Use conv2d(temperature, kernel, padding: 1) to apply the update across the grid.
Use a mask with which to preserve fixed heat sources or walls.

Jacobi iteration with conv2d

Implements the four-neighbor heat-update rule as a convolution and applies fixed boundary values with an elementwise conditional.

let kernel = MLXArray(converting: [
    0, 0.25, 0,
    0.25, 0, 0.25,
    0, 0.25, 0,
]).reshaped(1, 3, 3, 1)

var temperature = heatSources

let next = conv2d(temperature, kernel, padding: 1)
temperature = which(heatMask, heatSources, next)

Faster convergence and automatic differentiation

Jacobi iteration is easy to compute but can require about N² iterations for a grid with side length N. Successive Over-Relaxation uses the same convolution kernel, an omega parameter, and in-place-style updates to converge faster; MLX Swift simulates the in-place effect by updating red and black checkerboard cells in alternating phases.

The curve-fitting example demonstrates MLX Swift's grad transformation. A quadratic function and mean-squared-error loss are written normally, then grad(loss) produces a function that computes the gradient with respect to the parameter vector. The optimization loop updates parameters with gradient descent and calls eval each step.

SOR uses ω to overshoot updates and correct over iterations.
Red/black checkerboard masks let newly updated neighbors participate in the next half-step.
grad derives gradients automatically, avoiding hand-written derivative code.
MLX also includes optimizers such as SGD, Adam, and RMSprop for more general optimization workflows.

Successive Over-Relaxation update

Uses omega and checkerboard masks to approximate in-place SOR updates in an array-oriented MLX computation.

let ω: Float = 2.0 / (1.0 + sin(Float.pi / Float(max(M, N))))
let redMask = checkerboard(rows: M, cols: N, phase: 0)
let blackMask = checkerboard(rows: M, cols: N, phase: 1)

let sorRed = ω * conv2d(temperature, kernel, padding: 1) + (1 - ω) * temperature
temperature = which(redMask, sorRed, temperature)
temperature = which(heatMask, heatSources, temperature)

let sorBlack = ω * conv2d(temperature, kernel, padding: 1) + (1 - ω) * temperature
temperature = which(blackMask, sorBlack, temperature)
temperature = which(heatMask, heatSources, temperature)

Curve fitting with automatic differentiation

Defines a model and loss, transforms the loss into a gradient function, and runs gradient descent over the parameter array.

func f(_ θ: MLXArray) -> MLXArray {
    θ[0] + θ[1] * x + θ[2] * x ** 2
}

func loss(_ θ: MLXArray) -> MLXArray {
    mean((f(θ) - y) ** 2)
}

var θ = zeros([numParams])
let gradLoss = grad(loss)

for _ in 0 ..< steps {
    let g = gradLoss(θ)
    θ = θ - learningRate * g
    eval(θ)
}

Toolkit and ecosystem

Beyond the examples, MLX includes linear algebra, FFTs, n-dimensional convolutions, reductions, scans, indexing, random number generation, and more. The session also points to packages built on MLX Swift, including the core mlx-swift framework, mlx-swift-lm for Swift language-model implementations, and mlx-swift-examples for example applications.

Example areas mentioned include LLM integration, stable diffusion, model training and fine-tuning, and the session's numerical examples. MLX's shared concepts across Swift, Python, C++, and C make it practical to prototype in one front end and transfer patterns to another.

Start with the mlx-swift repository for framework documentation and tests.
Use mlx-swift-examples for complete example applications.
Look at mlx-swift-lm for Swift language-model implementations.
Python-side projects such as mlx-lm and mlx-vlm show additional ecosystem work.

Explore numerical computing in Swift with MLX