MLX

快速开始 | 安装 |
文档 |
示例

MLX 是苹果机器学习研究团队推出的，专为 Apple Silicon 设计的机器学习数组框架。

MLX 的主要特性包括：

• 熟悉的 API：MLX 提供了与 NumPy 高度相似的 Python API。同时，它还拥有功能完整的 C++、C 和 Swift API，这些 API 与 Python API 高度一致。MLX 还提供了像 mlx.nn 和 mlx.optimizers 这样的高级包，其 API 设计紧跟 PyTorch，以简化复杂模型的构建。

• 可组合的函数变换：MLX 支持可组合的函数变换，用于自动微分、自动向量化和计算图优化。

• 惰性计算：MLX 中的计算是惰性的。数组仅在需要时才被具体化。

• 动态图构建：MLX 中的计算图是动态构建的。改变函数参数的形状不会触发耗时的编译过程，调试也简单直观。

• 多设备支持：操作可以在任何支持的设备上运行（目前包括 CPU 和 GPU）。

• 统一内存：MLX 与其他框架的一个显著区别在于其统一内存模型。MLX 中的数组位于共享内存中。对 MLX 数组的操作可以在任何支持的设备类型上执行，而无需传输数据。

MLX 由机器学习研究者为机器学习研究者设计。该框架旨在易于使用，同时又能高效地训练和部署模型。框架本身的设计在概念上也力求简洁。我们希望研究者能够轻松地扩展和改进 MLX，以便快速探索新想法。

MLX 的设计灵感来源于 NumPy、PyTorch、Jax 和 ArrayFire 等框架。

示例

MLX 示例仓库 提供了丰富的示例，包括：

• Transformer 语言模型 训练。
• 使用 LLaMA 进行大规模文本生成，以及使用 LoRA 进行微调。
• 使用 Stable Diffusion 生成图像。
• 使用 OpenAI 的 Whisper 进行语音识别。

快速开始

请参阅文档中的 快速开始指南。

安装

MLX 已发布在 PyPI 上。要在 macOS 上安装 MLX，请运行：

pip install mlx

要在 Linux 上安装 CUDA 后端，请运行：

pip install mlx[cuda]

要安装仅支持 CPU 的 Linux 包，请运行：

pip install mlx[cpu]

有关从源代码构建 C++ 和 Python API 的更多信息，请查阅 安装文档。

贡献

请查看 贡献指南 以获取有关为 MLX 做出贡献的更多信息。有关从源代码构建和运行测试的详细信息，请参阅 文档。

我们感谢所有的 贡献者。如果您为 MLX 做出了贡献并希望被致谢，请在您的拉取请求中将您的名字添加到列表中。

引用 MLX

MLX 软件套件最初由 Awni Hannun、Jagrit Digani、Angelos Katharopoulos 和 Ronan Collobert 平等贡献开发。如果您在研究中发现 MLX 有用并希望引用它，请使用以下 BibTex 条目：

@software{mlx2023,
  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
  url = {https://github.com/ml-explore},
  version = {0.0},
  year = {2023},
}