🦀 mlx-rs v0.1.0 — Train Neural Nets on Apple Silicon, in Rust

🦀 mlx-rs v0.1.0 — Train Neural Nets on Apple Silicon, in Rust

The first release of mlx-rs — a batteries-included Rust interface to Apple's [MLX](https://github.com/ml-explore/mlx) framework. Build, train, and run ML models natively on M-series chips with Rust's safety guarantees.

This isn't just array bindings — it's a full training stack.

What's Inside

Core

• Safe Array type with automatic memory management via Drop, lazy evaluation, and unified CPU/GPU memory
• Dynamic FFI generation via bindgen — works with any mlx-c version on your machine
• Full linear algebra suite (SVD, QR, Cholesky, eigendecomposition, solve) and FFT operations

Neural Networks

• Module / ModuleParams trait system with a custom #[derive(ModuleParams)] proc macro — annotate fields with #[param], #[module], and #[state] and the boilerplate writes itself
• Layers: Linear, Conv1d/Conv2d, Embedding, LSTM, GRU, RNN, LayerNorm, RMSNorm, Dropout, pooling, Flatten, Sequential
• Full Transformer stack: MultiHeadAttention, TransformerEncoder, TransformerDecoder with KV cache
• Activations: ReLU, GELU, Sigmoid, Tanh, Softmax, LogSoftmax

Training

• Automatic differentiation: grad, value_and_grad, vjp, jvp
• JIT compilation and gradient checkpointing
• 8 optimizers: SGD (± momentum), Adam, AdamW, AdaGrad, RMSprop, Lion, Adafactor
• 7 loss functions: MSE, cross-entropy, BCE, L1, Smooth L1, KL divergence, cosine embedding

I/O & Quantization

• SafeTensors and NumPy save/load
• Quantize/dequantize operations (affine, configurable bit-width and group size)
• Distributed primitives: all_sum, all_gather, send, recv

Examples

Train a CNN on synthetic MNIST in ~30 lines of Rust. Also included: MLP, Transformer, RNN, Mistral inference, FFT, quantization, and distributed demos.

let model = RefCell::new(Sequential::new(vec![
    Box::new(Conv2d::new(1, 16, [3,3], [2,2], [1,1], [1,1], 1, true, &key)?),
    Box::new(ReLU::new()),
    Box::new(Conv2d::new(16, 32, [3,3], [2,2], [1,1], [1,1], 1, true, &key)?),
    Box::new(ReLU::new()),
    Box::new(Flatten::new()),
    Box::new(Linear::new(1568, 10, true, &key)?),
]));

let (loss, grads) = value_and_grad(|p| { /* forward + loss */ }, &params)?;
optimizer.update(params.iter_mut().collect(), grads)?;

Getting Started

Requires macOS with Apple Silicon, CMake, and a local build of [mlx-c](https://github.com/ml-explore/mlx-c). See the [README](https://github.com/MisterEkole/mlx-rs#-installation--setup) for full setup instructions.

Status

⚠️ Early release. APIs will evolve. Contributions, issues, and feedback are very welcome.

---

MIT Licensed. Not affiliated with Apple.