MojoGP

Active development: MojoGP is pre-1.0 software. Expect sharp edges, API changes, incomplete routes.

MojoGP is a Python-first exact Gaussian Process regression library backed by JIT-compiled Mojo GPU kernels. It exists to make exact GP training practical without depending on Torch: the runtime package depends on NumPy, SymPy, tqdm, and Mojo/MAX.

Why MojoGP

1. No torch dependency: NumPy, SymPy, tqdm, and Mojo/MAX only.
2. Materialized and matrix-free routes: dense materialized kernels and matrix-free kernel matvecs.
3. JIT-compiled GP models: models JIT-compile and kernels cached before training/prediction.
4. Multi-output GPs: ICM-style and LMC-style.
5. Discrete/categorical kernels: initial support for mixed continuous-categorical kernels.
6. Support for NVIDIA GPUs - support for AMD and Apple Silicon is on the roadmap.

Features

Mixed means continuous plus categorical/discrete inputs.

Examples

See notebooks/examples/ for runnable examples covering:

• single-output GPs
• multi-output workflows
• predictive uncertainty
• categorical variables
• observation-noise variants
• posterior sampling
• model persistence

Install

MojoGP is compatible with Python 3.10 and 3.11.

Install the PyPI package for your NVIDIA GPU:

pip install "mojogp[sm89]"

Use PyPI for package installs and Target for building from source.

Pinned install:

pip install "mojogp[sm89]==VERSION"

Build from Source

To build locally:

git clone https://github.com/caspbian/mojogp.git
cd mojogp
pip install -e .
task build GPU_TARGET=sm_89

Replace sm_89 with the Target value that matches your GPU. The build creates the complete split native engine set used by training, prediction, mixed, multi-output, and LMC routes.

Hello World

import numpy as np
from mojogp import RBF, SingleOutputGP

rng = np.random.default_rng(0)
X = np.linspace(-3, 3, 2000, dtype=np.float32).reshape(-1, 1)
y = (np.sin(2.0 * X[:, 0]) + 0.05 * rng.standard_normal(len(X))).astype(np.float32)

gp = SingleOutputGP(RBF())
gp.fit(X, y, max_iterations=50, method="matrix_free")

X_test = np.linspace(-4, 4, 128, dtype=np.float32).reshape(-1, 1)
mean, std = gp.predict(X_test, return_std=True, variance_method="love")

References

Bonilla, E.V., Chai, K. and Williams, C. (2007). Multi-task Gaussian Process Prediction. [online] Neural Information Processing Systems. Available at: https://papers.nips.cc/paper_files/paper/2007/hash/66368270ffd51418ec58bd793f2d9b1b-Abstract.html.

Bruinsma, W.P., Perim, E., Tebbutt, W., Scott, H.J., Solin, A. and Turner, R.E. (2019). Scalable Exact Inference in Multi-Output Gaussian Processes. [online] arXiv.org. Available at: https://arxiv.org/abs/1911.06287 [Accessed 22 May 2026].

Charlier, B., Feydy, J., Glaunès, J.A., Collin, F.-D. and Durif, G. (2020). Kernel Operations on the GPU, with Autodiff, without Memory Overflows. [online] arXiv.org. Available at: https://arxiv.org/abs/2004.11127 [Accessed 22 May 2026].

Chen, T., Huber, C., Lin, E. and Zaid, H. (2026). Preconditioning without a preconditioner using randomized block Krylov subspace methods. ETNA - Electronic Transactions on Numerical Analysis, [online] 65, pp.63–92. doi:https://doi.org/10.1553/etna_vol65s63.

Dong, K., Eriksson, D., Nickisch, H., Bindel, D. and Wilson, A.G. (2017). Scalable Log Determinants for Gaussian Process Kernel Learning. [online] arXiv.org. Available at: https://arxiv.org/abs/1711.03481 [Accessed 22 May 2026].

Gardner, J.R., Pleiss, G., Bindel, D., Weinberger, K.Q. and Wilson, A.G. (2021). GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration. arXiv:1809.11165 [cs, stat]. [online] Available at: https://arxiv.org/abs/1809.11165.

Godoy, W., Melnichenko, T., Valero-Lara, P., Elwasif, W., Fackler, P., Ferreira Da Silva, R., Teranishi, K. and Vetter, J. (2025). Mojo: MLIR-based Performance-Portable HPC Science Kernels on GPUs for the Python Ecosystem. Proceedings of the SC ’25 Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis, [online] pp.2114–2128. doi:https://doi.org/10.1145/3731599.3767573.

Harbrecht, H., Peters, M. and Schneider, R. (2012). On the low-rank approximation by the pivoted Cholesky decomposition. Applied numerical mathematics, 62(4), pp.428–440. doi:https://doi.org/10.1016/j.apnum.2011.10.001.

Perez, R.C., Veiga, D. and Garnier, J. (2025). A reproducible comparative study of categorical kernels for Gaussian process regression, with new clustering-based nested kernels. [online] arXiv.org. Available at: https://arxiv.org/abs/2510.01840 [Accessed 22 May 2026].

Peter, Wu, H. and Wu, C.Y. (2008). Gaussian Process Models for Computer Experiments With Qualitative and Quantitative Factors. 50(3), pp.383–396. doi:https://doi.org/10.1198/004017008000000262.

Pleiss, G., Gardner, J.R., Weinberger, K.Q. and Wilson, A.G. (2018). Constant-Time Predictive Distributions for Gaussian Processes. [online] arXiv.org. Available at: https://arxiv.org/abs/1803.06058 [Accessed 22 May 2026].

Rakitsch, B., Lippert, C., Borgwardt, K. and Stegle, O. (2026). It is all in the noise: Efficient multi-task Gaussian process inference with structured residuals. Advances in Neural Information Processing Systems, [online] 26. Available at: https://proceedings.neurips.cc/paper/2013/hash/59c33016884a62116be975a9bb8257e3-Abstract.html [Accessed 22 May 2026].

Rasmussen, C.E. and Williams, C.K.I. (2008). Gaussian processes for machine learning. Cambridge, Mass. Mit Press.

Roustant, O., Padonou, E., Deville, Y., Clément, A., Perrin, G., Giorla, J. and Wynn, H. (2018). Group kernels for Gaussian process metamodels with categorical inputs. [online] arXiv.org. Available at: https://arxiv.org/abs/1802.02368 [Accessed 22 May 2026].

Saves, P., Diouane, Y., Bartoli, N., Lefebvre, T. and Morlier, J. (2023). A mixed-categorical correlation kernel for Gaussian process. Neurocomputing, [online] 550, p.126472. doi:https://doi.org/10.1016/j.neucom.2023.126472.

Shashanka Ubaru, Chen, J. and Saad, Y. (2017). Fast Estimation of $tr(f(A))$ via Stochastic Lanczos Quadrature. SIAM Journal on Matrix Analysis and Applications, 38(4), pp.1075–1099. doi:https://doi.org/10.1137/16m1104974.

Wilson, A.G. and Nickisch, H. (2026). Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP). [online] arXiv.org. Available at: https://arxiv.org/abs/1503.01057 [Accessed 22 May 2026].

Wilson, J.T., Borovitskiy, V., Terenin, A., Mostowsky, P. and Deisenroth, M.P. (2020). Pathwise Conditioning of Gaussian Processes. [online] arXiv.org. Available at: https://arxiv.org/abs/2011.04026 [Accessed 22 May 2026].

Zhou, Q., Peter Z.G. Qian and Zhou, S. (2011). A Simple Approach to Emulation for Computer Models With Qualitative and Quantitative Factors. Technometrics, 53(3), pp.266–273. doi:https://doi.org/10.1198/tech.2011.10025.

License

MIT