Repository files navigation

Quickstart |Installation |Documentation |Examples

MLX is an array framework for machine learning on Apple silicon,brought to you by Apple machine learning research.

Some key features of MLX include:

• Familiar APIs: MLX has a Python API that closely follows NumPy. MLXalso has fully featured C++,C, andSwift APIs, which closely mirrorthe Python API. MLX has higher-level packages likemlx.nn andmlx.optimizers with APIs that closely follow PyTorch to simplify buildingmore complex models.

• Composable function transformations: MLX supports composable functiontransformations for automatic differentiation, automatic vectorization,and computation graph optimization.

• Lazy computation: Computations in MLX are lazy. Arrays are onlymaterialized when needed.

• Dynamic graph construction: Computation graphs in MLX are constructeddynamically. Changing the shapes of function arguments does not triggerslow compilations, and debugging is simple and intuitive.

• Multi-device: Operations can run on any of the supported devices(currently the CPU and the GPU).

• Unified memory: A notable difference from MLX and other frameworksis theunified memory model. Arrays in MLX live in shared memory.Operations on MLX arrays can be performed on any of the supporteddevice types without transferring data.

MLX is designed by machine learning researchers for machine learningresearchers. The framework is intended to be user-friendly, but still efficientto train and deploy models. The design of the framework itself is alsoconceptually simple. We intend to make it easy for researchers to extend andimprove MLX with the goal of quickly exploring new ideas.

The design of MLX is inspired by frameworks likeNumPy,PyTorch,Jax, andArrayFire.

TheMLX examples repo has avariety of examples, including:

• Transformer language model training.
• Large-scale text generation withLLaMA andfinetuning withLoRA.
• Generating images withStable Diffusion.
• Speech recognition withOpenAI's Whisper.

See thequick startguidein the documentation.

MLX is available onPyPI. To install the Python API, run:

Withpip:

pip install mlx

Withconda:

conda install -c conda-forge mlx

Checkout thedocumentationfor more information on building the C++ and Python APIs from source.

Check out thecontribution guidelines for more informationon contributing to MLX. See thedocs for moreinformation on building from source, and running tests.

We are grateful for all ofourcontributors. If you contributeto MLX and wish to be acknowledged, please add your name to the list in yourpull request.

The MLX software suite was initially developed with equal contribution by AwniHannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you findMLX useful in your research and wish to cite it, please use the followingBibTex entry:

@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},}