mojolearn

GPU machine learning in Mojo, on hardware the originals cannot reach.

mojolearn ports the GPU implementations that established today's tree and classical algorithms, CatBoost, cuVS, cuML and RAFT, into one Mojo source that runs on Apple silicon through Metal and, from the same source, on NVIDIA through CUDA and AMD through HIP. Every estimator has two numeric modes. FAST is the upstream's shipped behavior. IDENTICAL produces bit-identical models on every supported GPU vendor, with a per-stage certificate that proves it rather than a hash that hopes so.

The scikit-learn shapes are kept. The defaults follow the upstream each algorithm mirrors, and every place that differs from scikit-learn is named on the class.

Install in five minutes

Requirements are the lowest that can run it, not the machine it was built on. Any Apple silicon Mac from the M1 up, macOS 11 or later, Python 3.10 through 3.14, numpy 1.24 or later. The wheel and every runtime library in it are built for macOS 11 and the M1 instruction set, and the build refuses anything newer. There is no CPU path; every estimator runs on the GPU.

python3 -m venv .venv && source .venv/bin/activate
pip install mojolearn

import numpy as np
import mojolearn

rng = np.random.default_rng(0)
X = rng.random((100_000, 20), dtype=np.float32)
y = (X[:, 0] + X[:, 1] > 1.0).astype(np.float32)

gb = mojolearn.GradientBoosting(loss="Logloss", n_estimators=200, max_depth=6)
gb.fit(X, y)
print(gb.predict_proba(X[:5]))

rf = mojolearn.RandomForestClassifier(n_estimators=100, max_depth=12, random_state=7)
print(rf.fit(X, y).predict(X[:5]))

km = mojolearn.KMeans(n_clusters=8).fit(X)
print(km.cluster_centers_.shape, mojolearn.numeric_mode())

To train the bit-identical way, set one environment variable before the import. Nothing is rebuilt; the wheel carries both binary sets.

MOJOLEARN_NUMERIC_MODE=identical python train.py

mojolearn.numeric_mode() reports the mode that actually loaded, read back from the binary, so a run cannot be mislabeled by accident.

What is in 0.1

Estimators save to and load from .npz files, and a model fitted on a Mac loads and predicts identically on an NVIDIA or AMD box (95 of 95 models, probabilities included, in the E2 certificate below).

Not on the surface in 0.1, named rather than omitted: MultiClassOneVsAll from Python, Intel and Qualcomm GPU columns, and a CPU fallback of any kind.

The two numeric modes

FAST is the default. Histograms flush through float atomics, library reductions follow the hardware warp width, and the last bits of a model can move between two runs on the same device. That is CatBoost's shipped behavior and it is the fastest the hardware goes.

IDENTICAL pins every pathway that can move a bit. The enumeration of those pathways is IDENTITY_PATHS.md, 32 rows, each with what IDENTICAL does about it and the status of that closure. The mode makes exactly three kinds of move, pin a machine-derived parameter to a frozen floor, replace an order-dependent operation with a fixed-point or fixed-shape one, or refuse by name. There is no fourth move and no "usually fine".

Identity coverage, measured on three vendors

One source tree, one commit, byte-identical inputs proven by hash, and the device plus toolchain as the only variable. Every fit writes an identity trace card, one hash per training stage, and cards are diffed stage by stage.

The Logloss row in E1 diverged on a device exp and log seam that the ledger had named before any hardware ran, and round 2 of E2 closed that class for the whole matrix. The unsupervised estimators have their own ledger rows and Apple to AMD cards (UNSUPERVISED_IDENTITY.md).

Reproduce a certificate with the runbook in E1_RUNBOOK.md. On the Mac, one fit and its card is

pixi run python tools/e1_traced_fit.py --fit et_clf --out cards/mac

and the comparison against a card from another box is python tools/identity_trace_diff.py cards/mac/et_clf.card cards/box/et_clf.card.

Hardware

Support is one source; validation is what the certificates say and nothing more. Performance has been measured on exactly one machine, an M4 laptop with 10 cores and 16 GB, and the numbers below carry that.

Benchmarks

Method first. We run NVIDIA's gbm-bench harness unmodified except for registering our arms and letting its imports survive a machine with no CUDA. Their timing code, their datasets, their metrics, their parameters for their arms. On Apple silicon none of the established libraries has a working GPU path, so the comparison is our Metal path against their CPU path, on the same machine, all arms of one invocation interleaved in one process because the box drifts across thermal windows. Every arm records a sha256 of its prediction vector beside its timing, so determinism is visible in the results themselves. The reproduction command is

pixi run -e gbmbench bash bench/external/run_gbm_bench.sh year 500 gbdt
pixi run -e gbmbench bash bench/external/run_gbm_bench.sh covtype 100 forest

and the results and the machine record land in bench/results/gbm_bench_*. These are the last interleaved runs of 2026-08-22 on the M4 described above.

Six interleaved year runs over the day ranged from 1.14x to 1.68x as the laptop warmed; the table shows the last one, not the best one. The CatBoost pair compares symmetric trees only, because LightGBM has no symmetric mode. The higgs GBDT accuracy gap is stated because it is there.

Limitations and refusals

• GPU only. No CPU fallback exists and none is planned for this release.
• The wheel is macOS arm64. CUDA and HIP are source builds.
• Validated on one M4 for performance; correctness and identity validated on the M4, an H100, an MI325X and an MI300X through the certificates above.
• IDENTICAL refuses rather than guessing. A GPU column that misses the frozen identity floor, a k-NN arm that needs a warp primitive a column does not have, or k > 256 on the identical k-NN selector raise with a named reason instead of returning a model that might not match.
• KMeans.n_init is 1. GradientBoosting defaults are CatBoost's.
• No Intel or Qualcomm GPU column has hardware to build for yet.
• Numerics are float32 on the device; no float64 device path.

Provenance and licensing

Every line of Mojo here was written in this repository and Andrew Hendel holds its copyright. What it mirrors is the design of the upstreams, file for file where the toolchain allows, under the rule copy, do not improve, so that a port can be checked against its original. That makes the ported directories derivative works under Apache-2.0 section 4, and NOTICE carries each upstream's attribution. PORTED_MAP.tsv maps every ported file to its origin and status. This is not a clean-room reimplementation and must not be described as one.

License Apache-2.0, in LICENSE. AUTHORS.md records who wrote what. Not affiliated with, endorsed by, or sponsored by YANDEX LLC, NVIDIA, Meta, or Modular, Inc. MAX and Mojo are trademarks of Modular, Inc. used under license.

Citing

If mojolearn is useful in work you publish, please cite it. CITATION.cff is what GitHub's Cite this repository button reads; each GitHub release is archived on Zenodo with its own DOI.

Where things are

• IDENTITY_PATHS.md, the bit-identity ledger.
• E1_RESULTS.md, E2_RESULTS.md, E1_RUNBOOK.md, the cross-vendor certificates and how to reproduce them.
• PORTING_RULES.md, PORTING.md, how the port is done and every numbered deviation from the upstream.
• VENDOR_LIBS.md, where a vendor library is called instead of hand-written.
• docs/DESIGN_NOTES.md, the working notes that were this README during the port.
• docs/PYPI_RELEASE.md, how a release is built, verified and published.
• ROADMAP.md, what comes next.