MetalML: run scikit-learn on your Mac's GPU with one import change

GitHub: https://github.com/tudor-22/MetalML

MetalML started with a practical idea: keep the familiar scikit-learn workflow while making supported numerical operations available on a Mac's GPU. Instead of introducing a separate training framework, it exposes estimators with the fit, predict and transform methods Python users already know.

# pip install metalml
import numpy as np
import metalml as ml  # was: import sklearn as ml

X = np.random.default_rng(7).normal(size=(5000, 32)).astype(np.float32)
y = X[:, 0] * 3 - X[:, 1]
model = ml.linear_model.Ridge(alpha=1.0).fit(X, y)
print(model.predict(X[:5]))
print(model.operation_backends_)  # which ops ran on Metal vs CPU

Status

MetalML is experimental. GPU execution is verified on an Intel Mac with a Radeon Pro 5300M; Apple Silicon validation and tuning are next. If you have an M-series Mac, benchmark reports and issues are very welcome.

What runs on the GPU?

Acceleration depends on the estimator and the operation. Scalers accelerate statistics and normalisation, Ridge uses GPU matrix products with a CPU solve, eligible KMeans configurations accelerate Lloyd iterations and distances, and trees, forests and kernel SVMs accelerate prediction while fitting stays on the CPU. One pipeline can mix CPU and GPU steps.

Precision and fallback by design

GPU operations use dense float32. Sparse, float64 and integer inputs keep scikit-learn's path, so precision never changes silently. Anything MetalML doesn't implement is forwarded to your installed scikit-learn, and importing MetalML never patches sklearn itself.

Inspect instead of assuming a speedup

operation_backends_ and fallback_reasons_ show exactly what ran where, and ml.config_context(backend="cpu") lets you compare.
⭐ Star it, try it, or open an issue on GitHub: https://github.com/tudor-22/MetalML