MLX

An open-source array framework from Apple machine learning research, first published in December 2023 under the MIT License. It is built so researchers can train and run Machine Learning (ML) models on Apple silicon, using a Python Application Programming Interface (API) that closely follows NumPy, with matching C++, C, and Swift bindings. Arrays live in the unified memory that M-series chips share between the Central Processing Unit (CPU) and the Graphics Processing Unit (GPU), so an operation can run on either device without copying tensors between separate memories. A later CUDA backend extended the same API to NVIDIA GPUs on Linux, and a CPU-only Linux package covers machines with no accelerator.

The design is inspired by NumPy, PyTorch, JAX, and ArrayFire. Higher-level packages mlx.nn and mlx.optimizers follow PyTorch closely enough that neural-network modules port with little rewriting. The project is hosted under Apple’s ml-explore GitHub organization and was initially developed with equal contribution by Awni Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert.

Unified Memory and Device Placement

On a typical CUDA stack, a tensor lives in host memory or on a discrete GPU, and the programmer copies it before compute starts. Apple silicon does not split those pools. The CPU and GPU already address the same physical memory, and the framework is written around that fact rather than around .to("cuda")-style transfers.

• No per-array device: An array is created without a location argument. The device is chosen when the operation runs, as a stream such as mx.cpu or mx.gpu.
• Parallel streams: Independent operations can run on the CPU and GPU at the same time. When one operation depends on a result produced on the other device, the scheduler inserts the wait so the consumer does not start early.
• Workload split: Apple’s documentation timed a large matrix multiply on the GPU plus a long sequence of tiny elementwise ops on the CPU at about 1.4 ms on an M1 Max, versus about 2.8 ms with everything on the GPU.
• GPU, not the Neural Engine: Documented device types are CPU and GPU (Metal on Apple silicon). The framework does not expose Apple’s Neural Processing Unit (NPU) as a first-class device. Neural Engine inference remains the domain of Core ML.
• Swift on Apple platforms: The Swift API runs the same runtime on Mac, iPhone, iPad, and Vision Pro, not only in Python on a laptop.

Lazy Evaluation and Function Transforms

Calling an operation records a compute graph. Values are materialized only when the graph is evaluated: an explicit mx.eval(), printing an array, converting it to NumPy, saving it, or reading a scalar with .item(). Unused branches can skip real compute. A large mlx.nn.Module can be constructed and then given float16 weights without paying for a full float32 initialization, because construction itself does no work until evaluation.

Graphs are built dynamically, so changing input shapes does not force a slow recompile the way a static XLA graph can. mx.compile() is available when you do want to fuse a graph for repeated execution. Transforms compose: each one returns a function that can be transformed again.

• Automatic differentiation: mx.grad and mx.value_and_grad operate on functions, not on tensors marked requires_grad. There is no backward() or zero_grad().
• Automatic vectorization: mx.vmap batches a function over extra axes without writing the batching loops by hand.
• Compilation: mx.compile fuses the recorded graph so hot training or inference loops spend less time on Python and kernel-launch overhead.

Language Models and Fine-Tuning

mlx-lm is the Python package for loading, generating, quantizing, and fine-tuning Large Language Models (LLMs) on top of the array runtime. It is the usual entry point for running a local model on a Mac without a cloud GPU, and Apple has since shown the same stack for on-device agentic workflows and for sharding large models across several Macs.

• Hugging Face Hub: Thousands of converted models ship through the mlx-community organization. The package exposes mlx_lm.generate and mlx_lm.chat on the command-line interface (CLI), plus a Python API (load, generate, stream_generate).
• Quantization: mlx_lm.convert can quantize a Hub checkpoint (commonly 4-bit) and optionally upload the result. Converted weights are stored in MLX format, typically safetensors, not GGUF. Runtimes such as Ollama sit on llama.cpp and a different file format.
• Fine-tuning: Low-rank adaptation (LoRA), QLoRA, and full-model updates are supported, including on quantized bases. That is Parameter-Efficient Fine-Tuning (PEFT) in practice: a small adapter instead of a second copy of the base model.
• Long context: A rotating key-value cache, prompt caching, and a configurable prefill step size keep peak memory in check on long prompts without re-computing the prefix every turn.
• Distributed: mx.distributed covers multi-machine inference and training. Later releases added RDMA over Thunderbolt through the JACCL backend on recent macOS versions.

A companion examples repository covers speech recognition (Whisper), image generation (Stable Diffusion, FLUX), encoder models (BERT), and Mixture of Experts (MoE) language models, among others.

Backends and Install Targets

The Python package is on PyPI. macOS wheels require Apple silicon, a native arm Python 3.10 or newer (Rosetta or x86 Python will not match a wheel), and macOS 14.0 or later. Linux extra packages add CUDA and CPU-only builds.

Backend	Platform	Accelerator	Install
Metal	macOS 14+ on Apple silicon	GPU in unified memory, shared with the CPU	pip install mlx
CUDA 12 or 13	Linux, NVIDIA SM 7.5+	Discrete NVIDIA GPU (not unified memory)	pip install mlx with the cuda12 or cuda13 extra
CPU	Linux, glibc 2.35+	Host CPU (BLAS)	pip install mlx with the cpu extra