Hi all,
I’ve been building Mojolearn, a machine learning library written in Mojo with a Python API. The idea is to run the same supported AI workload on Apple, NVIDIA or AMD GPUs or Intel, AMD and Apple CPUs and get exactly the same results, down to the bits.
The library covers classical ML, including random forests, Extra Trees and gradient boosting, alongside neural training tools for MLPs, Transformers and Mamba, including hybrid Mamba / attention models like Samba.
Why care? If two machines give the same bits, a result is reproducible in a strict sense. You can run inference on different devices and get the exact same output on both. You can replay a training run for a reviewer, move work between supported hardware, or find exactly where two runs diverged, with no tolerance to argue about. That matters in regulated work like finance, healthcare and government, and in research where you need to explain why a result changed.
There’s a Mac benefit too. Random forests, Extra Trees and gradient boosting can train on the Apple GPU. The trees also have a separate fast mode (not bitwise identical) that beat scikit-learn on some Mac workloads. It depends on the model, data and settings. For random forests and extra trees, the training savings are significant.
Mojo lets one GPU source compile for all three vendors. Getting identical bits still meant controlling the order and form of every calculation.
This is an alpha. Identity is verified for the specific configurations in the support matrix, and it can cost speed, especially for neural networks. The library is future proofed for Intel and Qualcomm if Mojo ever supports those GPUs.
Have you needed to reproduce a model on different hardware? What would you try this on?
Install on a supported platform: pip install mojolearn
GitHub, setup and supported configurations: