mojolearn

GPU machine learning in Mojo for Apple, NVIDIA and AMD GPUs, with an explicit reproducibility contract.

CatBoost, XGBoost, LightGBM and cuML have no Metal backend. mojolearn ports the GPU kernels of CatBoost, cuML, cuVS, RAFT and FAISS into one Mojo source that builds for Apple Metal, NVIDIA CUDA and AMD HIP, under the repository's copy-do-not-improve rule (PORTING_RULES.md), so GPU tree training and GPU classical learning run on Apple silicon, hardware the originals cannot reach. On top of that port sits a numerical contract. The same machine-learning workload can produce different bits on different GPUs, changing predictions, learned models and subsequent training updates. In mojolearn's opt-in identical mode, supported inference and training return the same bits on certified Apple, NVIDIA and AMD GPUs. The claim is proven by stage-level identity cards and separating sabotage tests, never inferred from a final-output hash, and it applies only to configurations recorded in the support matrix.

What the contract covers today, with the evidence each claim rests on:

• Neural inference and training. Mamba and transformer forward computations agree bit for bit across the three vendors on their recorded fixtures, as do gradients, optimizer updates and checkpoint bytes in fixed-shape transformer training. A two-block, 34,944-parameter byte-level language model trained on real text ran 128 steps with byte-identical parameters, gradients, optimizer state and loss on Apple Metal, NVIDIA CUDA and AMD HIP; held-out loss fell from 5.5413 to 2.8436 on all three (three-vendor record). Checkpoint continuation between NVIDIA and AMD, in both directions, preserves the uninterrupted training trajectory (cross-vendor record). Metal checkpoint resume remains open.
• Trees and classical learning. Gradient boosting, random forests, Extra Trees, k-means, DBSCAN, k-NN, PCA, truncated SVD, OLS, ridge, logistic regression, FP32 matrix multiplication, isolation forest and ARIMA filtering carry three-vendor cards for recorded configurations. Model state and recorded training stages match, not only predictions.
• UMAP. Neighbor selection and iterative updates match across the three vendors on named fixtures.

Two other modes sit beside identical, selectable at runtime on every supported estimator:

Bitwise identity carries implementation and execution costs. Measured against cuML, cuBLAS and PyTorch, identical mode is competitive on some measured tree workloads and substantially slower on many classical, matrix and neural workloads. Those measurements reflect both the numerical constraints and optimization gaps in the current kernels. The numbers are in the accompanying paper; the raw records behind them live under bench/results/.

The 0.6.1 source in this checkout defaults to identical mode. Set MOJOLEARN_NUMERIC_MODE=fast or deterministic explicitly to opt into another supported mode. Certification remains configuration-specific; this default does not certify every feature. The published 0.6.0 files are unchanged.

Who this is for

• People on Apple silicon who want GPU gradient boosting, random forests, Extra Trees, clustering, nearest neighbors, decompositions and linear models without leaving the machine.
• People who need a reproducibility contract on trees or classical models, same bits on repeated runs or across vendors, and will pay for it in time. The cost is small on some measured tree workloads and large elsewhere; see the paper before deciding.
• Not yet people training real neural networks. The certified trainers are fixed small shapes, an MLP and the two-block byte LM above. Larger models, other shapes and other optimizers are outside the evidence, and the byte-LM native trainer is not in any published wheel.

Install

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

Version 0.6.0 is published on PyPI as an alpha API release for macOS arm64 and AMD Linux x86-64 (gfx942). It exposes public linalg, umap, training, Mamba and Transformer APIs, including UMAP transform/CSR support. NVIDIA Linux remains source-build-only. The Linux wheel contains HIP binaries, not CUDA. New byte-LM native training needs a separate source build. The wheels preserve identified base native/runtime bytes; newer Python API exposure does not inherit every numerical certificate. See the release evidence and support matrix for exact artifacts and limits. There is no CPU fallback. Run the diagnostic command before depending on a new machine:

mojolearn doctor

The exact wheel, architecture, Python, and evidence boundaries live in SUPPORT_MATRIX.md. Source builds may support hardware outside the architectures packaged in a released wheel; that is not the same as released-wheel support.

Project status

Stability and release cadence

mojolearn went from 0.1.0 on 2026-08-23 to 0.6.0 on 2026-09-06, six PyPI releases in two weeks (0.1.0, 0.2.0, 0.3.0, 0.3.1, 0.5.0, 0.6.0; 0.3.2 and 0.4.0 are recorded in CHANGELOG.md but were not published to PyPI). One release was yanked. 0.3.0, published 2026-08-30 as the first release with a Linux wheel, had been compiled for the build machine's CPU and carried unconditional AVX-512 instructions in its host code, so every numeric mode died with SIGILL on any x86-64 host without AVX-512. It is yanked on PyPI with the reason "SIGILL on x86-64 without AVX-512; use 0.3.1". 0.3.1 pinned the Linux baseline to x86-64-v3 and added a gate on the shipped binary; the defect and both gates are documented in packaging/linux/isa_baseline_linux.py and packaging/wheel_ci.py.

The Python API is beta and will change between minor versions. The stable surface is the set of numerical profiles (fast, deterministic, identical) and the certified configurations recorded in SUPPORT_MATRIX.md: a profile version changes only through an explicit decision, and a numerical change must either prove itself bit-inert or introduce a new profile version. For production or archival work pin both the package version and the numeric profile, in code or through MOJOLEARN_NUMERIC_MODE. A certificate names a commit, a configuration (the fixture, the numeric profile, the parameters) and the devices it ran on, and never more. A newer version, a different shape or an unrun vendor column is not covered by it.

Maintenance and bus factor

The project has one maintainer today. Three things limit what that means for a reader.

Every claim in this repository is backed by a recorded artifact under bench/results/ that names its commit, device, toolchain, mode and limitations, and each is reproducible from the commands in the docs (verification, conformance bundles, release runbook). Historical cards and investigations under bench/results/ and archive/ are evidence, not current guidance; SUPPORT_MATRIX.md is updated only from recorded evidence.

Contributions are governed by CONTRIBUTING.md and GOVERNANCE.md. A contributor needs one GPU of any vendor and marks the vendor columns they did not run cross-vendor-pending; closing a cross-vendor claim is a maintainer job. Any change that can move identical bits must show that it is bit-inert, supply a separating fixture and a profile-version decision, or add a named refusal. External pull requests get an admission report and a hosted CPU report; there is no GPU automation and no automatic merge. Governance uses lazy consensus with a seven-day objection window, maintainership is explicitly transferable, a sole maintainer records nominations in a public issue, and the succession steps for a sole maintainer (nominate two successors, transfer access, document release and certification steps, rotate credentials, publish open blockers) are written down. The code is Apache-2.0.

You can verify a certificate without trusting the maintainer. On any supported GPU, MOJOLEARN_NUMERIC_MODE=identical python -m mojolearn verify runs a pinned fixture, captures its stage-level identity card and compares it with the reference card shipped in the installation; python -m mojolearn check-fixture checks the fixture's input hashes without a GPU. Recorded cards carry stage tags, dtypes, element counts and raw-bit hashes and are compared with tools/identity_trace_diff.py, the one comparator the repository uses. python -m mojolearn conformance exports and validates bundles so another implementation can compare itself without running Mojo, and tools/verify_umap_qualification.py rechecks retained release evidence against a wheel without GPU work. One local run establishes one build on one device; a cross-vendor claim needs every named leg, and the cards for each leg are in the tree.

Quick start

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)

model = mojolearn.GradientBoosting(
    loss="Logloss", n_estimators=200, max_depth=6,
    numeric_mode="deterministic",
)
model.fit(X, y)
print(model.predict_proba(X[:5]))
print(model.numeric_mode_used(), mojolearn.vendor())

Choose a process default with mojolearn.set_numeric_mode("identical"), or set the starting default before import:

MOJOLEARN_NUMERIC_MODE=identical python train.py

More than one tier may be loaded in one process through per-estimator numeric_mode= arguments.

Public API

Classical estimators include:

• Gradient boosting, random forests, and Extra Trees
• K-means, nearest-neighbor estimators, DBSCAN, hierarchical and spectral clustering
• PCA, truncated SVD, linear and logistic regression, ridge, lasso, and elastic net
• SVC, SVR, kernel density, isolation forest, and Gaussian-process regression
• Exponential smoothing and batched ARIMA
• UMAP embeddings with dense Euclidean input and 2D/3D spectral initialization; version 0.6.0 adds unseen-sample transformation and CSR graph storage

Additional modules provide scoring metrics, FP32 matrix multiplication, optimizer/training primitives, and reference-pinned Mamba and transformer blocks. These surfaces do not all have the same validation depth; consult the support matrix before treating an experimental surface as release-qualified.

UMAP in the 0.5.0 API supports fitting and embedding the supplied samples:

X = np.array([0, 1, 2.2, 4, 6.5, 10, 14.5, 20], dtype=np.float32)[:, None]
embedding = mojolearn.UMAP(
    n_neighbors=3, n_components=2, n_epochs=4, random_state=19,
    numeric_mode="identical",
).fit_transform(X)

The 0.5.0 implementation stores a dense graph and does not support transform. In 0.6.0, public fitting stores the graph in CSR form, using O(n_samples × n_neighbors) graph space, and transform(X_new) embeds unseen samples against a frozen fitted model. Input remains a dense Euclidean array; CSR describes internal graph storage. Exact neighbor search still performs quadratic pair comparisons.

Source checks for the integrated fit/transform API passed all three numeric modes on Apple, NVIDIA and AMD. The named IDENTICAL held-out embeddings match across all three vendors. The macOS 0.6.0 candidate also passed clean installed fit/transform and quality checks. See the version-specific evidence.

Transformation retains private training data and embedding copies. Changing parameters or numeric mode requires refitting, and changing query batching can change results. Supervised targets, alternate metrics and alternate initialization remain unsupported.

The APIs intentionally resemble scikit-learn, but mojolearn is not a drop-in replacement. Defaults follow the upstream GPU implementation mirrored by an algorithm where applicable. Unsupported parameters raise explicitly rather than being silently ignored.

What the identity claim means

Fix a source commit, a supported configuration, a seed and byte-identical input. On any two certified machines, every recorded training stage has the same bits, and either model produces exactly the same predictions. This is a claim about the trained model, not byte-for-byte equality of archive metadata. If a configuration cannot meet the contract, the library raises a named error instead of silently returning a possibly different model; a refusal is reported as a refusal, never counted as a pass.

Cross-vendor identity is a profile, not a statement that every GPU operation is universally identical. A profile fixes relevant reduction order, partitioning, FMA policy, flush-to-zero seams, transcendental spellings, and tie rules. A numerical change must either prove bit-inertness against the profile or introduce a new profile version. Additional devices must pass the same identity checks; the guarantee covers only devices and configurations that have.

The project distinguishes four artifact classes:

source check -> Python binding -> built native artifact -> installed wheel

Evidence for one class does not automatically validate the next. Current certificates, configurations, and outstanding vendor legs are listed in SUPPORT_MATRIX.md. Historical cards and investigations under bench/results/ and archive/ are evidence, not current guidance.

Limitations

What will get in your way first:

• GPU hardware is required. There is no CPU fallback, and the library refuses rather than silently running elsewhere.
• mojolearn is not a drop-in replacement for scikit-learn, CatBoost or cuML. Parameter coverage is intentionally smaller than any of them, and unsupported parameters raise.
• Source builds need the Mojo toolchain through pixi, and one build targets one GPU architecture. NVIDIA Linux is source-build-only today.
• The support matrix is honest about gaps. Several public surfaces still have vendor legs or independent-reference checks pending, and an unrun column is pending, never inferred.

And the standing limits of the contract itself:

• Released-wheel support is narrower than source-build support.
• fast deliberately makes no repeatability promise.
• deterministic does not promise agreement between different devices.
• identical covers certified profiles and fixtures, not arbitrary untested shapes or future toolchains.
• Some recent Python and neural-operator surfaces still have vendor legs or independent-reference checks pending.
• Parameter coverage is intentionally smaller than scikit-learn, CatBoost, or cuML.
• The experimental k-NN selector remains behind an explicit build flag; normal wheel builds retain the existing dispatch.

mojolearn is beta software. Pin the package version and numerical profile for production or archival work.

Development

Start with docs/START_HERE.md. The shortest full local check is pixi run probe.

A numerical test counts as evidence only after a separating arm demonstrates that it fails when the relevant rule is broken. Contributors need one supported GPU; maintainers close cross-vendor certification columns.

Current priorities are in ROADMAP.md. See also verification, release, porting rules, contributing, governance, and attribution.

Provenance and citation

The shipped implementation is Mojo. Algorithmic designs derive in part from CatBoost, cuML, cuVS, RAFT, and FAISS; exact provenance and licenses are in NOTICE, DERIVATION_MAP.tsv, source headers, and the archived derivation ledger.

To cite mojolearn, use CITATION.cff. The concept DOI is 10.5281/zenodo.22068632.