Explore numerical computing in Swift with MLX
Swift & Data 进阶 14m

使用 MLX 探索 Swift 中的数值计算

Explore numerical computing in Swift with MLX

2026年6月10日

在 Apple 官方观看视频

一句话判断

这场 Session 最值得关注的一件事是：MLX Swift 让 Swift 真正拥有了 NumPy 级别的矩阵运算和自动微分能力，写复杂数学公式终于不用和底层 Metal 死磕了。

这场 Session 讲了什么

以前在 Apple 平台搞数值计算或端侧 AI 训练，体验非常割裂。要么用 Accelerate 和 Metal Performance Shaders (MPS) 写一堆晦涩的底层代码，要么在 Python 里训好模型再费劲转成 CoreML。Swift 原生缺乏一套能优雅表达高维数学运算的工具。

现在 MLX Swift 把 n 维数组 (n-dimensional arrays)、惰性求值 (lazy evaluation) 和自动微分 (automatic differentiation) 原生搬进了 Swift。你写的代码看起来就像数学公式，比如 z = z * z + c，底层会自动构建计算图并扔到 GPU 上并行执行。

这对搞端侧模型微调、科学计算和复杂图像处理的开发者是巨大利好。你终于可以在一个类型安全、有完整 Xcode 调试体验的环境里，直接写训练循环和物理模拟，彻底干掉跨语言调用的胶水代码。

值得深挖的点

惰性求值 (Lazy Evaluation) 的取舍

MLX Swift 没有采用 PyTorch 那种即时执行 (eager execution) 模式，而是选择了惰性求值。当你对数组做加减乘除时，它不会立刻计算，而是默默构建一个计算图 (compute graph)。直到你调用 eval() 或者打印结果时，它才会把整个图发给 GPU 执行。

Apple 这么做的理由很充分：GPU 启动 kernel 的开销很大，把一堆小操作融合 (operator fusion) 成一个大 kernel 能榨干 Apple Silicon 的性能。在 Mandelbrot 和热分布的例子中，这种设计让代码极其简洁，且轻松获得 10 倍以上的加速。

代价是调试体验的妥协。你没法在循环里随便 print 中间变量，因为此时计算还没发生。如果你在循环里忘了写 eval()，计算图会无限膨胀，直接 OOM (Out of Memory)。这要求开发者在脑海里建立明确的“图构建”与“图执行”的边界。

自动微分 (Automatic Differentiation) 降维打击

grad() 函数的引入是这场 Session 最大的杀手锏。以前要在 Swift 里做梯度下降，你得自己手推导数公式，或者用 MPS 手写反向传播的 Shader。现在，你只需要定义前向的 loss 函数，套一层 grad(loss)，MLX 就能自动帮你算出梯度。

这直接打通了端侧模型训练和复杂曲线拟合的任督二脉。相比于 Swift for TensorFlow 当年过于激进的编译器魔改，MLX Swift 采用库级别的函数转换，更轻量也更稳定。只要你的函数是用 MLX 数组操作写的，梯度就能自动推导，这让 Swift 在科学计算领域的表达力直接追平了 JAX。

代码片段

1. Mandelbrot 分形计算

场景：用数组级运算替代双重 for 循环，让 GPU 并行处理整个复数网格，代码量锐减且速度提升 10 倍。

import MLX

// 构建复数网格 c = x + iy
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² + c
    z = z * z + c                       
    // 统计未逃逸的迭代次数
    counts = counts + (abs(z) .< 2)     
}

坑：纯 Swift 写标量循环很容易，但换成 MLX 后，千万不要在数组操作里混用原生的 Float 标量运算，必须全用 MLX 的张量操作符。

2. 自动微分曲线拟合

场景：利用 grad 自动推导二次函数的梯度，无需手写任何导数公式即可完成参数优化。

// 定义前向预测函数
func f(_ θ: MLXArray) -> MLXArray {
    θ[0] + θ[1] * x + θ[2] * x ** 2
}

// 定义均方误差损失
func loss(_ θ: MLXArray) -> MLXArray {
    mean((f(θ) - y) ** 2)
}

var θ = zeros([numParams])
// 魔法发生在这里：自动获取 loss 对 θ 的梯度函数
let gradLoss = grad(loss)

for _ in 0 ..< steps {
    let g = gradLoss(θ)         
    // 梯度下降更新参数
    θ = θ - learningRate * g    
    // 强制求值并清空计算图，防止内存爆炸
    eval(θ)                     
}

坑：循环末尾的 eval(θ) 绝对不能省。它不仅触发 GPU 计算，还会截断计算图，漏写会导致内存随迭代次数线性增长。

最佳实践

  • 已有项目的迁移策略：如果你的项目只是做简单的向量加减或矩阵乘法，继续用 Accelerate，它的 CPU 优化依然无敌。如果是纯推理，CoreML 仍是首选。MLX Swift 最适合需要动态控制流、自定义训练循环或复杂物理模拟的场景。
  • 新项目的采用建议：从 mlx-swift-examples 抄代码起步。利用 Python 端丰富的生态（如 mlx-lm）做算法原型验证，然后用 MLX Swift 重写核心逻辑发布到 iOS/macOS，两者底层 C++ 逻辑完全一致，迁移成本极低。
  • 实战中容易踩的坑：在实现类似热传导的物理模拟时，尽量用 conv2d 配合掩码 (mask) 来替代手写的邻居遍历。Session 中展示的 SOR (逐次超松弛法) 结合红黑棋盘格掩码，用几行代码就把收敛速度提升了 100 倍，这是典型的“用算子思维代替循环思维”。

还有什么值得关注

  • 多语言前端共享底层：MLX 支持 Swift、Python、C++ 和 C，这意味着你可以用 Python 跑通算法，直接无缝翻译成 Swift 上架，彻底告别 CoreML 转换的玄学报错。
  • 现成的 LLM 生态：mlx-swift-lm 仓库已经提供了 Swift 原生的语言模型实现，可以直接拿来做端侧大模型集成和微调，不用自己从头写 Transformer 算子。
MLX Swift Numerical Computing AutoDiff GPU