Apple has open sourced MLX, an array framework for machine learning on Apple silicon (i.e your laptop). 

Developed by Apple’s machine learning research team, MLX introduces a range of features tailored to meet the demands of researchers, ensuring a streamlined experience for model training and deployment.

MLX comes equipped with several noteworthy features:

• Familiar APIs: MLX’s Python API closely aligns with NumPy, while the fully-featured C++ API mirrors the Python version. Additionally, higher-level packages such as mlx.nn and mlx.optimizers simplify model building by adhering to PyTorch conventions.

• Composable Function Transformations: MLX introduces composable function transformations, enabling automatic differentiation, vectorization, and computation graph optimization.

• Lazy Computation: Computation in MLX is designed to be lazy, ensuring that arrays are only materialized when necessary, optimizing computational efficiency.

• Dynamic Graph Construction: MLX adopts dynamic graph construction, eliminating slow compilations triggered by changes in function argument shapes. This approach simplifies the debugging process.

• Multi-Device Support: MLX allows operations to seamlessly run on supported devices, including the CPU and GPU, providing flexibility for developers.

• Unified Memory Model: MLX introduces a unified memory model, deviating from other frameworks. Arrays reside in shared memory, enabling operations on MLX arrays across different device types without data movement.

Drawing inspiration from established frameworks like NumPy, PyTorch, Jax, and ArrayFire, MLX combines key features to create a robust and versatile platform.

The MLX examples repository showcases the framework’s capabilities, including transformer language model training, large-scale text generation, image generation with Stable Diffusion, and speech recognition using OpenAI’s Whisper.

MLX is conveniently available on PyPi, and installation of the Python API is a straightforward process with the command: pip install mlx.