MLX Array Framework is a high-performance array library for machine learning on Apple silicon. Designed by Apple machine learning researchers, it provides a unified memory model that allows operations to run seamlessly across CPU and GPU without data copying.
- Familiar APIs: Features a Python API closely following NumPy, alongside fully featured C++, C, and Swift APIs. Includes
mlx.nnandmlx.optimizerspackages inspired by PyTorch. - Composable Function Transformations: Supports automatic differentiation, vectorization, and computation graph optimization.
- Lazy Computation: Arrays are only materialized when strictly necessary, optimizing memory usage.
- Dynamic Graph Construction: No slow compilations when changing input shapes; debugging is intuitive.
- Unified Memory Model: Arrays live in shared memory, eliminating the need for data transfers between devices.
To get started, install the package via pip:
bash pip install mlx
Basic array operations:
python import mlx.core as mx
a = mx.array([1, 2, 3]) b = mx.array([4, 5, 6])
result = a + b
print(mx.eval(result))
MLX requires macOS 13.3+ and an Apple Silicon machine (M1/M2/M3).
For detailed installation instructions, including building from source, please visit the official documentation.
The MLX examples repository contains a variety of demonstrations:
- LLMs: Text generation with LLaMA and finetuning with LoRA.
- Vision: Image generation with Stable Diffusion.
- Audio: Speech recognition with OpenAI's Whisper.
- Training: Transformer language model training.
Full documentation is available at https://ml-explore.github.io/mlx/.
MLX is an open-source project intended to be useful for machine learning researchers. We welcome contributions that align with the project's goal of simplicity and efficiency.
This project is licensed under the BSD 3-Clause License.