MLX 是由 Apple machine learning research 推出的类 NumPy 数组框架,面向 Apple silicon 上高效、灵活的机器学习。它的 Python API 接近 NumPy,同时提供与 Python API closely follows 的完整 C++ API。MLX 的设计受到 PyTorch、JAX 和 ArrayFire 等框架启发,但采用统一内存模型:数组位于共享内存中,在支持的设备类型之间执行操作时无需进行数据复制。目前官网列出的支持设备类型为 CPU 和 GPU,文档同时覆盖安装、使用、示例、Python API、C++ API 以及扩展开发等内容。
核心功能
类 NumPy 数组与多语言 API
MLX 提供类 NumPy 的数组框架,Python API 遵循 NumPy 的接口方式,但存在少数差异。除 Python API 外,MLX 还提供功能完整的 C++ API,且 C++ API 与 Python API 保持接近。官网文档列出了数组、数据类型、操作、随机数、线性代数、FFT、神经网络、优化器和内存管理等 API 参考内容。
可组合的函数变换
MLX 支持可组合的函数变换,用于自动微分、自动向量化和计算图优化。这些变换可以组合使用,官网将其列为 MLX 与 NumPy 之间的主要差异之一。相关文档还提供了 Function Transforms 与 Compilation 主题,覆盖函数变换和编译能力。
惰性计算
MLX 的计算采用惰性方式执行,数组只会在需要时被具体化。官网将 Lazy computation 作为框架的核心设计,并提供 Lazy Evaluation 使用指南。该文档体系还包括保存和加载数组、数组索引、数值精度以及环境变量等相关内容。
多设备执行与统一内存
MLX 的操作可以运行在受支持的设备上,包括 CPU 和 GPU。其统一内存模型让 MLX 数组位于共享内存中,在不同支持设备类型之间执行数组操作时无需进行数据复制。官网还提供 Devices and Streams、Unified Memory、Using Streams 和 Memory Management 等使用与 API 文档。
机器学习与大模型示例
官网提供 Linear Regression、Multi-Layer Perceptron、LLM inference、Data Parallelism 和 Tensor Parallelism 等示例,覆盖线性回归、多层感知器、大语言模型推理、数据并行和张量并行。文档还提供 Writing a Fast KV Cache 指南,说明 MLX 的使用内容包含机器学习模型计算与大语言模型推理相关场景。
扩展、导出与分布式能力
MLX 文档包含 Distributed Communication、Exporting Functions 以及 Export Functions,用于分布式通信和函数导出;同时提供 Custom Extensions in MLX、Custom Metal Kernels、Metal Debugger、Metal Logging 和 Using MLX in C++ 等进一步阅读内容。Python API 参考中还列出了 Metal 与 CUDA 页面。
使用方式
官网提供 Build and Install 安装入口,并以 Quick Start Guide 作为入门路径。使用者可以通过 Python API 或 C++ API 调用 MLX,官方文档覆盖数组操作、函数变换、编译、保存与加载数组,以及与 NumPy 和其他框架之间的转换。官网正文没有说明是否需要注册、是否需要自备模型 Key,也没有提供 Web 端、浏览器插件或独立客户端信息。
适用人群与场景
需要在 Apple silicon 上进行机器学习开发的 Python 使用者,可以使用接近 NumPy 的 API 处理数组与模型计算;使用 C++ 的开发者可以参考完整的 C++ API 和 Using MLX in C++ 文档。进行线性回归、多层感知器实验的机器学习学习者,可使用官网对应示例;需要开展 LLM inference、数据并行或张量并行工作的开发者,可参考官网提供的相关示例与 KV Cache 指南。希望研究 Metal 扩展或自定义 Metal kernel 的开发者,则可查看进一步阅读中的扩展开发内容。
注意事项
官网当前列出的支持设备类型为 CPU 和 GPU,产品定位为面向 Apple silicon 的机器学习框架。正文没有提供价格、免费额度、注册要求或数据处理政策等信息。