MLX — MLX 0.32.0 documentation

MLX — 面向 Apple Silicon 的高效机器学习框架

简介

MLX 是一个由 Apple 机器学习研究团队推出的 NumPy 风格数组框架，专为 Apple Silicon 设备上的高效、灵活机器学习而设计。它提供与 NumPy 高度相似的 Python API，同时也拥有功能完备的 C++ API，后者的设计紧密遵循 Python API。MLX 的设计灵感来源于 PyTorch、Jax 和 ArrayFire 等框架，但其核心差异在于统一内存模型 —— 数组存储在共享内存中，可在 CPU 和 GPU 等不同设备上执行操作而无需数据拷贝。

主要功能

• 可组合函数变换：支持自动微分、自动向量化以及计算图优化等可组合的函数变换。
• 惰性计算：MLX 中的计算是惰性的，数组仅在需要时才被物化（materialized）。
• 多设备支持：操作可在任何支持的设备类型（CPU、GPU 等）上运行，当前支持 CPU 和 GPU。
• 统一内存模型：数组位于共享内存，可在不同设备类型间操作而无需显式数据拷贝。
• 丰富的 API：涵盖数组操作、数据类型、设备与流、导出函数、随机数、线性代数、FFT、神经网络、优化器、分布式通信等，并提供 Python 和 C++ 两种 API。
• 与 NumPy 及其他框架的互操作：支持数组与 NumPy 及其他框架之间的转换与导出。
• 编译与优化：支持函数编译、计算图优化等高级功能。

适用场景

• 在 Apple Silicon 设备上进行机器学习研究与开发：利用统一内存模型和惰性计算，可高效实现模型训练和推理。
• 需要自动微分和向量化的科学计算：可组合的函数变换简化了梯度计算和批量操作。
• 多设备并行计算：支持数据并行和模型并行（如张量并行），适合分布式机器学习场景。
• 快速原型开发：Python API 与 NumPy 类似，易于上手，且提供丰富的示例（如线性回归、多层感知机、LLM 推理等）。
• 需要底层控制的高级用户：C++ API 和自定义 Metal 内核功能允许深度优化和扩展。