mojogp 0.26.6.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (127) hide show
  1. mojogp-0.26.6.0/LICENSE +21 -0
  2. mojogp-0.26.6.0/MANIFEST.in +15 -0
  3. mojogp-0.26.6.0/PKG-INFO +222 -0
  4. mojogp-0.26.6.0/README.md +148 -0
  5. mojogp-0.26.6.0/mojogp/__init__.py +162 -0
  6. mojogp-0.26.6.0/mojogp/_multi_output_backend.py +223 -0
  7. mojogp-0.26.6.0/mojogp/_provider_lifecycle.py +275 -0
  8. mojogp-0.26.6.0/mojogp/_provider_runtime.py +258 -0
  9. mojogp-0.26.6.0/mojogp/_routes.py +28 -0
  10. mojogp-0.26.6.0/mojogp/_version.py +3 -0
  11. mojogp-0.26.6.0/mojogp/benchmark/__init__.py +0 -0
  12. mojogp-0.26.6.0/mojogp/benchmark/config.py +93 -0
  13. mojogp-0.26.6.0/mojogp/codegen_engine/__init__.py +254 -0
  14. mojogp-0.26.6.0/mojogp/codegen_engine/compiler.py +307 -0
  15. mojogp-0.26.6.0/mojogp/codegen_engine/differentiation.py +53 -0
  16. mojogp-0.26.6.0/mojogp/codegen_engine/emit/__init__.py +20 -0
  17. mojogp-0.26.6.0/mojogp/codegen_engine/emit/builder.py +75 -0
  18. mojogp-0.26.6.0/mojogp/codegen_engine/emit/fn_ptr_module.py +1211 -0
  19. mojogp-0.26.6.0/mojogp/codegen_engine/emit/gpu_kernel.py +1666 -0
  20. mojogp-0.26.6.0/mojogp/codegen_engine/emit/module.py +360 -0
  21. mojogp-0.26.6.0/mojogp/codegen_engine/emit/mojo_printer.py +134 -0
  22. mojogp-0.26.6.0/mojogp/codegen_engine/expressions.py +650 -0
  23. mojogp-0.26.6.0/mojogp/codegen_engine/ir.py +294 -0
  24. mojogp-0.26.6.0/mojogp/codegen_engine/overrides.py +80 -0
  25. mojogp-0.26.6.0/mojogp/codegen_engine/passes/__init__.py +35 -0
  26. mojogp-0.26.6.0/mojogp/codegen_engine/passes/cse.py +94 -0
  27. mojogp-0.26.6.0/mojogp/codegen_engine/passes/dead_code.py +73 -0
  28. mojogp-0.26.6.0/mojogp/codegen_engine/passes/forward_hoist.py +139 -0
  29. mojogp-0.26.6.0/mojogp/codegen_engine/passes/gradient_sharing.py +25 -0
  30. mojogp-0.26.6.0/mojogp/codegen_engine/passes/inv_ls.py +218 -0
  31. mojogp-0.26.6.0/mojogp/codegen_engine/passes/numerical_stability.py +112 -0
  32. mojogp-0.26.6.0/mojogp/codegen_engine/passes/simplify.py +17 -0
  33. mojogp-0.26.6.0/mojogp/codegen_engine/passes/strength.py +83 -0
  34. mojogp-0.26.6.0/mojogp/codegen_engine/registry.py +27 -0
  35. mojogp-0.26.6.0/mojogp/codegen_engine/schedule.py +144 -0
  36. mojogp-0.26.6.0/mojogp/compiler.py +202 -0
  37. mojogp-0.26.6.0/mojogp/feature_support.py +902 -0
  38. mojogp-0.26.6.0/mojogp/gp.py +3949 -0
  39. mojogp-0.26.6.0/mojogp/kernel.py +1635 -0
  40. mojogp-0.26.6.0/mojogp/kernels/__init__.mojo +9 -0
  41. mojogp-0.26.6.0/mojogp/kernels/bbmm_gpu_kernels.mojo +356 -0
  42. mojogp-0.26.6.0/mojogp/kernels/bbmm_types.mojo +530 -0
  43. mojogp-0.26.6.0/mojogp/kernels/categorical_composable.mojo +409 -0
  44. mojogp-0.26.6.0/mojogp/kernels/categorical_kernel.mojo +712 -0
  45. mojogp-0.26.6.0/mojogp/kernels/categorical_state.mojo +364 -0
  46. mojogp-0.26.6.0/mojogp/kernels/cg_solver.mojo +985 -0
  47. mojogp-0.26.6.0/mojogp/kernels/cg_with_provider.mojo +199 -0
  48. mojogp-0.26.6.0/mojogp/kernels/combined_inv_quad_logdet.mojo +1199 -0
  49. mojogp-0.26.6.0/mojogp/kernels/composable_kernel.mojo +1772 -0
  50. mojogp-0.26.6.0/mojogp/kernels/composite_matvec.mojo +1051 -0
  51. mojogp-0.26.6.0/mojogp/kernels/composite_prediction.mojo +1033 -0
  52. mojogp-0.26.6.0/mojogp/kernels/composite_provider.mojo +969 -0
  53. mojogp-0.26.6.0/mojogp/kernels/constants.mojo +32 -0
  54. mojogp-0.26.6.0/mojogp/kernels/cross_covariance_provider.mojo +457 -0
  55. mojogp-0.26.6.0/mojogp/kernels/gemm_matvec.mojo +272 -0
  56. mojogp-0.26.6.0/mojogp/kernels/gemm_tiled.mojo +158 -0
  57. mojogp-0.26.6.0/mojogp/kernels/generic_matvec.mojo +1641 -0
  58. mojogp-0.26.6.0/mojogp/kernels/gpu_utils.mojo +363 -0
  59. mojogp-0.26.6.0/mojogp/kernels/gradient_provider.mojo +1022 -0
  60. mojogp-0.26.6.0/mojogp/kernels/jit/__init__.mojo +5 -0
  61. mojogp-0.26.6.0/mojogp/kernels/jit/erased_provider.mojo +860 -0
  62. mojogp-0.26.6.0/mojogp/kernels/jit/fused_kronecker_provider.mojo +363 -0
  63. mojogp-0.26.6.0/mojogp/kernels/jit/jit_categorical_params.mojo +89 -0
  64. mojogp-0.26.6.0/mojogp/kernels/jit/jit_engine_binding_helpers.mojo +64 -0
  65. mojogp-0.26.6.0/mojogp/kernels/jit/jit_engine_bindings.mojo +1775 -0
  66. mojogp-0.26.6.0/mojogp/kernels/jit/jit_engine_lmc_bindings.mojo +1931 -0
  67. mojogp-0.26.6.0/mojogp/kernels/jit/jit_engine_multi_output_bindings.mojo +1547 -0
  68. mojogp-0.26.6.0/mojogp/kernels/jit/jit_lmc.mojo +1529 -0
  69. mojogp-0.26.6.0/mojogp/kernels/jit/jit_lmc_mixed.mojo +1557 -0
  70. mojogp-0.26.6.0/mojogp/kernels/jit/jit_mixed.mojo +1132 -0
  71. mojogp-0.26.6.0/mojogp/kernels/jit/jit_multi_output.mojo +817 -0
  72. mojogp-0.26.6.0/mojogp/kernels/jit/jit_multi_output_mixed.mojo +1272 -0
  73. mojogp-0.26.6.0/mojogp/kernels/jit/jit_prediction.mojo +2084 -0
  74. mojogp-0.26.6.0/mojogp/kernels/jit/jit_progress.mojo +65 -0
  75. mojogp-0.26.6.0/mojogp/kernels/jit/jit_training.mojo +1008 -0
  76. mojogp-0.26.6.0/mojogp/kernels/kernel_functions.mojo +1903 -0
  77. mojogp-0.26.6.0/mojogp/kernels/kernel_materialization.mojo +372 -0
  78. mojogp-0.26.6.0/mojogp/kernels/kernel_params.mojo +166 -0
  79. mojogp-0.26.6.0/mojogp/kernels/kronecker_direct_provider.mojo +730 -0
  80. mojogp-0.26.6.0/mojogp/kernels/kronecker_gpu_kernels.mojo +112 -0
  81. mojogp-0.26.6.0/mojogp/kernels/kronecker_preconditioner.mojo +822 -0
  82. mojogp-0.26.6.0/mojogp/kernels/kronecker_provider.mojo +453 -0
  83. mojogp-0.26.6.0/mojogp/kernels/lanczos.mojo +634 -0
  84. mojogp-0.26.6.0/mojogp/kernels/lanczos_with_provider.mojo +246 -0
  85. mojogp-0.26.6.0/mojogp/kernels/lmc_bbmm_step.mojo +369 -0
  86. mojogp-0.26.6.0/mojogp/kernels/lmc_bbmm_step_mf.mojo +334 -0
  87. mojogp-0.26.6.0/mojogp/kernels/lmc_composite_provider.mojo +512 -0
  88. mojogp-0.26.6.0/mojogp/kernels/lmc_preconditioner.mojo +719 -0
  89. mojogp-0.26.6.0/mojogp/kernels/lmc_provider.mojo +1438 -0
  90. mojogp-0.26.6.0/mojogp/kernels/love_provider.mojo +659 -0
  91. mojogp-0.26.6.0/mojogp/kernels/matvec_provider.mojo +1464 -0
  92. mojogp-0.26.6.0/mojogp/kernels/mixed_composite_matvec.mojo +1121 -0
  93. mojogp-0.26.6.0/mojogp/kernels/mixed_composite_provider.mojo +837 -0
  94. mojogp-0.26.6.0/mojogp/kernels/mixed_composite_training.mojo +717 -0
  95. mojogp-0.26.6.0/mojogp/kernels/multi_output_prediction.mojo +1511 -0
  96. mojogp-0.26.6.0/mojogp/kernels/native_numerics.mojo +901 -0
  97. mojogp-0.26.6.0/mojogp/kernels/pivoted_cholesky.mojo +2444 -0
  98. mojogp-0.26.6.0/mojogp/kernels/preconditioner_trait.mojo +82 -0
  99. mojogp-0.26.6.0/mojogp/kernels/profiling_config.mojo +4 -0
  100. mojogp-0.26.6.0/mojogp/kernels/py_conversion.mojo +108 -0
  101. mojogp-0.26.6.0/mojogp/kernels/task_covariance.mojo +1070 -0
  102. mojogp-0.26.6.0/mojogp/kernels/training.mojo +2053 -0
  103. mojogp-0.26.6.0/mojogp/kernels/training_types.mojo +538 -0
  104. mojogp-0.26.6.0/mojogp/kernels/training_utils.mojo +587 -0
  105. mojogp-0.26.6.0/mojogp/kernels/utils.mojo +339 -0
  106. mojogp-0.26.6.0/mojogp/loader.py +429 -0
  107. mojogp-0.26.6.0/mojogp/multi_output_gp.py +7028 -0
  108. mojogp-0.26.6.0/mojogp/pathwise_prior.py +593 -0
  109. mojogp-0.26.6.0/mojogp/progress.py +448 -0
  110. mojogp-0.26.6.0/mojogp/settings.py +74 -0
  111. mojogp-0.26.6.0/mojogp/specialization/__init__.py +49 -0
  112. mojogp-0.26.6.0/mojogp/specialization/benchmark.py +110 -0
  113. mojogp-0.26.6.0/mojogp/specialization/builtins.py +192 -0
  114. mojogp-0.26.6.0/mojogp/specialization/decision.py +79 -0
  115. mojogp-0.26.6.0/mojogp/specialization/descriptor.py +73 -0
  116. mojogp-0.26.6.0/mojogp/specialization/profile.py +73 -0
  117. mojogp-0.26.6.0/mojogp/specialization/registry.py +69 -0
  118. mojogp-0.26.6.0/mojogp/specialization/rules.py +23 -0
  119. mojogp-0.26.6.0/mojogp/specialization/translate.py +39 -0
  120. mojogp-0.26.6.0/mojogp/utils.py +173 -0
  121. mojogp-0.26.6.0/mojogp.egg-info/PKG-INFO +222 -0
  122. mojogp-0.26.6.0/mojogp.egg-info/SOURCES.txt +125 -0
  123. mojogp-0.26.6.0/mojogp.egg-info/dependency_links.txt +1 -0
  124. mojogp-0.26.6.0/mojogp.egg-info/requires.txt +69 -0
  125. mojogp-0.26.6.0/mojogp.egg-info/top_level.txt +1 -0
  126. mojogp-0.26.6.0/pyproject.toml +88 -0
  127. mojogp-0.26.6.0/setup.cfg +4 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 caspbian
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,15 @@
1
+ include README.md
2
+ include LICENSE
3
+ recursive-include mojogp *.py *.mojo
4
+ recursive-exclude * __pycache__ *.py[cod]
5
+ exclude mojogp_jit_engine.so
6
+ prune archive
7
+ prune benchmark_data_vast
8
+ prune isolated_engine_loads
9
+ prune max_runtime_libs
10
+ prune max_sdk_bundle
11
+ prune node_modules
12
+ prune ref
13
+ prune results
14
+ prune tests
15
+ prune tmp
@@ -0,0 +1,222 @@
1
+ Metadata-Version: 2.4
2
+ Name: mojogp
3
+ Version: 0.26.6.0
4
+ Summary: GPU-accelerated Gaussian Process library using Mojo kernels
5
+ Author: caspbian
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/caspbian/mojogp
8
+ Project-URL: Documentation, https://caspbian.github.io/mojogp/
9
+ Project-URL: Repository, https://github.com/caspbian/mojogp
10
+ Project-URL: Issues, https://github.com/caspbian/mojogp/issues
11
+ Keywords: gaussian-processes,mojo,gpu,machine-learning
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Requires-Python: <3.12,>=3.10
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: numpy
22
+ Requires-Dist: mojo==0.25.7.0
23
+ Requires-Dist: max==25.7.0
24
+ Requires-Dist: sympy>=1.12
25
+ Requires-Dist: tqdm
26
+ Provides-Extra: test
27
+ Requires-Dist: gpytorch; extra == "test"
28
+ Requires-Dist: linear_operator; extra == "test"
29
+ Requires-Dist: torch>=2.11.0; extra == "test"
30
+ Requires-Dist: pytest>=7.0; extra == "test"
31
+ Requires-Dist: scipy>=1.7.0; extra == "test"
32
+ Requires-Dist: matplotlib>=3.5.0; extra == "test"
33
+ Requires-Dist: scikit-learn; extra == "test"
34
+ Requires-Dist: properscoring>=0.1; extra == "test"
35
+ Requires-Dist: psutil>=5.8.0; extra == "test"
36
+ Provides-Extra: docs
37
+ Requires-Dist: griffe>=1.14.0; extra == "docs"
38
+ Requires-Dist: mkdocs-material>=9.6.0; extra == "docs"
39
+ Requires-Dist: mkdocstrings[python]>=0.30.0; extra == "docs"
40
+ Provides-Extra: notebooks
41
+ Requires-Dist: marimo>=0.23.1; extra == "notebooks"
42
+ Requires-Dist: matplotlib; extra == "notebooks"
43
+ Requires-Dist: scikit-learn; extra == "notebooks"
44
+ Requires-Dist: tqdm; extra == "notebooks"
45
+ Provides-Extra: sm75
46
+ Requires-Dist: mojogp-cuda-sm75==0.26.6.0; extra == "sm75"
47
+ Provides-Extra: sm-75
48
+ Requires-Dist: mojogp-cuda-sm75==0.26.6.0; extra == "sm-75"
49
+ Provides-Extra: sm80
50
+ Requires-Dist: mojogp-cuda-sm80==0.26.6.0; extra == "sm80"
51
+ Provides-Extra: sm-80
52
+ Requires-Dist: mojogp-cuda-sm80==0.26.6.0; extra == "sm-80"
53
+ Provides-Extra: sm86
54
+ Requires-Dist: mojogp-cuda-sm86==0.26.6.0; extra == "sm86"
55
+ Provides-Extra: sm-86
56
+ Requires-Dist: mojogp-cuda-sm86==0.26.6.0; extra == "sm-86"
57
+ Provides-Extra: sm89
58
+ Requires-Dist: mojogp-cuda-sm89==0.26.6.0; extra == "sm89"
59
+ Provides-Extra: sm-89
60
+ Requires-Dist: mojogp-cuda-sm89==0.26.6.0; extra == "sm-89"
61
+ Provides-Extra: sm90
62
+ Requires-Dist: mojogp-cuda-sm90==0.26.6.0; extra == "sm90"
63
+ Provides-Extra: sm-90
64
+ Requires-Dist: mojogp-cuda-sm90==0.26.6.0; extra == "sm-90"
65
+ Provides-Extra: sm100
66
+ Requires-Dist: mojogp-cuda-sm100==0.26.6.0; extra == "sm100"
67
+ Provides-Extra: sm-100
68
+ Requires-Dist: mojogp-cuda-sm100==0.26.6.0; extra == "sm-100"
69
+ Provides-Extra: sm120
70
+ Requires-Dist: mojogp-cuda-sm120==0.26.6.0; extra == "sm120"
71
+ Provides-Extra: sm-120
72
+ Requires-Dist: mojogp-cuda-sm120==0.26.6.0; extra == "sm-120"
73
+ Dynamic: license-file
74
+
75
+ # MojoGP
76
+
77
+ > **Active development:** MojoGP is pre-1.0 software. Expect sharp edges, API changes, incomplete routes.
78
+
79
+ 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.
80
+
81
+ ## Why MojoGP
82
+
83
+ 1. No torch dependency: NumPy, SymPy, tqdm, and Mojo/MAX only.
84
+ 2. Materialized and matrix-free routes: dense materialized kernels and matrix-free kernel matvecs.
85
+ 3. JIT-compiled GP models: models JIT-compile and kernels cached before training/prediction.
86
+ 4. Multi-output GPs: ICM-style and LMC-style.
87
+ 5. Discrete/categorical kernels: initial support for mixed continuous-categorical kernels.
88
+ 6. Support for NVIDIA GPUs - support for AMD and Apple Silicon is on the roadmap.
89
+
90
+ ## Features
91
+
92
+ Mixed means continuous plus categorical/discrete inputs.
93
+
94
+ | Feature | Single-output continuous | Single-output mixed | Multi-output ICM continuous | Multi-output ICM mixed | Multi-output LMC continuous | Multi-output LMC mixed |
95
+ |---|---|---|---|---|---|---|
96
+ | Materialized training | alpha | experimental | experimental | experimental | alpha | experimental |
97
+ | Matrix-free training | alpha | experimental | experimental | experimental | alpha | experimental |
98
+ | Mean-only prediction | alpha | experimental | experimental | experimental | alpha | experimental |
99
+ | Exact variance | alpha | experimental | experimental | experimental | alpha | experimental |
100
+ | LOVE variance | alpha | experimental | experimental | experimental | experimental | experimental |
101
+ | Heterogeneous latent kernels | n/a | n/a | n/a | n/a | alpha | experimental |
102
+ | Active dimensions | alpha | experimental | experimental | experimental | alpha | experimental |
103
+ | ARD lengthscales | alpha | experimental | experimental | in development | alpha | in development |
104
+ | Additive kernel composites | alpha | experimental | experimental | experimental | alpha | in development |
105
+ | Product kernel composites | alpha | experimental | experimental | experimental | alpha | experimental |
106
+ | Save / load | alpha | experimental | experimental | experimental | alpha | experimental |
107
+ | Learned homoskedastic noise | alpha | experimental | experimental | experimental | alpha | experimental |
108
+ | Fixed observation noise | alpha | in development | alpha | in development | alpha | in development |
109
+ | Learned heteroskedastic noise | alpha | in development | not started | not started | not started | not started |
110
+ | Grouped noise | alpha | in development | alpha | in development | unsupported | unsupported |
111
+ | Posterior sampling | alpha | experimental | experimental | experimental | alpha | experimental |
112
+
113
+ ## Examples
114
+
115
+ See [`notebooks/examples/`](notebooks/examples/) for runnable examples covering:
116
+
117
+ - single-output GPs
118
+ - multi-output workflows
119
+ - predictive uncertainty
120
+ - categorical variables
121
+ - observation-noise variants
122
+ - posterior sampling
123
+ - model persistence
124
+
125
+ ## Install
126
+
127
+ PyPI packages are in development. For now, install MojoGP from source.
128
+
129
+ Minimum runtime dependencies are Python 3.10 or 3.11, NumPy, SymPy, tqdm,
130
+ Mojo, and MAX. The current build has been tested with Mojo 0.25.7.0 and MAX
131
+ 25.7.0.
132
+
133
+
134
+ ## Build From Source
135
+
136
+ For now, build MojoGP from source. This workflow clones the source, installs the
137
+ Python package, and builds one JIT engine locally for the GPU target you choose:
138
+
139
+ ```bash
140
+ git clone https://github.com/caspbian/mojogp.git
141
+ cd mojogp
142
+ pip install -e .
143
+ mojo build mojogp/kernels/jit/jit_engine_bindings.mojo \
144
+ --emit shared-lib \
145
+ -I mojogp/ \
146
+ -o mojogp_jit_engine.so \
147
+ --target-accelerator sm_89
148
+ ```
149
+
150
+ Use the Mojo accelerator target that matches your GPU:
151
+
152
+
153
+ | GPU family | Target |
154
+ |---|---|
155
+ | T4 / RTX 20-series / Quadro RTX | `sm_75` |
156
+ | A100 / A30 | `sm_80` |
157
+ | A40 / A10 / A16 / A2 / RTX 30-series / RTX A-series | `sm_86` |
158
+ | L4 / L40 / L40S / RTX 40-series / RTX Ada | `sm_89` |
159
+ | GH200 / H100 / H200 | `sm_90` |
160
+ | B200 / GB200 | `sm_100` |
161
+ | RTX PRO Blackwell / GeForce RTX 50-series | `sm_120` |
162
+
163
+
164
+ ## Hello World
165
+
166
+ ```python
167
+ import numpy as np
168
+ from mojogp import RBF, SingleOutputGP
169
+
170
+ rng = np.random.default_rng(0)
171
+ X = np.linspace(-3, 3, 2000, dtype=np.float32).reshape(-1, 1)
172
+ y = (np.sin(2.0 * X[:, 0]) + 0.05 * rng.standard_normal(len(X))).astype(np.float32)
173
+
174
+ gp = SingleOutputGP(RBF())
175
+ gp.fit(X, y, max_iterations=50, method="matrix_free")
176
+
177
+ X_test = np.linspace(-4, 4, 128, dtype=np.float32).reshape(-1, 1)
178
+ mean, std = gp.predict(X_test, return_std=True, variance_method="love")
179
+ ```
180
+
181
+ ## References
182
+ 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.
183
+
184
+ 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].
185
+
186
+ 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].
187
+
188
+ 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.
189
+
190
+ 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].
191
+
192
+ 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.
193
+
194
+ 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.
195
+
196
+ 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.
197
+
198
+ 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].
199
+
200
+ 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.
201
+
202
+ 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].
203
+
204
+ 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].
205
+
206
+ Rasmussen, C.E. and Williams, C.K.I. (2008). Gaussian processes for machine learning. Cambridge, Mass. Mit Press.
207
+
208
+ 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].
209
+
210
+ 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.
211
+
212
+ 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.
213
+
214
+ 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].
215
+
216
+ 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].
217
+
218
+ 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.
219
+
220
+ ## License
221
+
222
+ MIT
@@ -0,0 +1,148 @@
1
+ # MojoGP
2
+
3
+ > **Active development:** MojoGP is pre-1.0 software. Expect sharp edges, API changes, incomplete routes.
4
+
5
+ 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.
6
+
7
+ ## Why MojoGP
8
+
9
+ 1. No torch dependency: NumPy, SymPy, tqdm, and Mojo/MAX only.
10
+ 2. Materialized and matrix-free routes: dense materialized kernels and matrix-free kernel matvecs.
11
+ 3. JIT-compiled GP models: models JIT-compile and kernels cached before training/prediction.
12
+ 4. Multi-output GPs: ICM-style and LMC-style.
13
+ 5. Discrete/categorical kernels: initial support for mixed continuous-categorical kernels.
14
+ 6. Support for NVIDIA GPUs - support for AMD and Apple Silicon is on the roadmap.
15
+
16
+ ## Features
17
+
18
+ Mixed means continuous plus categorical/discrete inputs.
19
+
20
+ | Feature | Single-output continuous | Single-output mixed | Multi-output ICM continuous | Multi-output ICM mixed | Multi-output LMC continuous | Multi-output LMC mixed |
21
+ |---|---|---|---|---|---|---|
22
+ | Materialized training | alpha | experimental | experimental | experimental | alpha | experimental |
23
+ | Matrix-free training | alpha | experimental | experimental | experimental | alpha | experimental |
24
+ | Mean-only prediction | alpha | experimental | experimental | experimental | alpha | experimental |
25
+ | Exact variance | alpha | experimental | experimental | experimental | alpha | experimental |
26
+ | LOVE variance | alpha | experimental | experimental | experimental | experimental | experimental |
27
+ | Heterogeneous latent kernels | n/a | n/a | n/a | n/a | alpha | experimental |
28
+ | Active dimensions | alpha | experimental | experimental | experimental | alpha | experimental |
29
+ | ARD lengthscales | alpha | experimental | experimental | in development | alpha | in development |
30
+ | Additive kernel composites | alpha | experimental | experimental | experimental | alpha | in development |
31
+ | Product kernel composites | alpha | experimental | experimental | experimental | alpha | experimental |
32
+ | Save / load | alpha | experimental | experimental | experimental | alpha | experimental |
33
+ | Learned homoskedastic noise | alpha | experimental | experimental | experimental | alpha | experimental |
34
+ | Fixed observation noise | alpha | in development | alpha | in development | alpha | in development |
35
+ | Learned heteroskedastic noise | alpha | in development | not started | not started | not started | not started |
36
+ | Grouped noise | alpha | in development | alpha | in development | unsupported | unsupported |
37
+ | Posterior sampling | alpha | experimental | experimental | experimental | alpha | experimental |
38
+
39
+ ## Examples
40
+
41
+ See [`notebooks/examples/`](notebooks/examples/) for runnable examples covering:
42
+
43
+ - single-output GPs
44
+ - multi-output workflows
45
+ - predictive uncertainty
46
+ - categorical variables
47
+ - observation-noise variants
48
+ - posterior sampling
49
+ - model persistence
50
+
51
+ ## Install
52
+
53
+ PyPI packages are in development. For now, install MojoGP from source.
54
+
55
+ Minimum runtime dependencies are Python 3.10 or 3.11, NumPy, SymPy, tqdm,
56
+ Mojo, and MAX. The current build has been tested with Mojo 0.25.7.0 and MAX
57
+ 25.7.0.
58
+
59
+
60
+ ## Build From Source
61
+
62
+ For now, build MojoGP from source. This workflow clones the source, installs the
63
+ Python package, and builds one JIT engine locally for the GPU target you choose:
64
+
65
+ ```bash
66
+ git clone https://github.com/caspbian/mojogp.git
67
+ cd mojogp
68
+ pip install -e .
69
+ mojo build mojogp/kernels/jit/jit_engine_bindings.mojo \
70
+ --emit shared-lib \
71
+ -I mojogp/ \
72
+ -o mojogp_jit_engine.so \
73
+ --target-accelerator sm_89
74
+ ```
75
+
76
+ Use the Mojo accelerator target that matches your GPU:
77
+
78
+
79
+ | GPU family | Target |
80
+ |---|---|
81
+ | T4 / RTX 20-series / Quadro RTX | `sm_75` |
82
+ | A100 / A30 | `sm_80` |
83
+ | A40 / A10 / A16 / A2 / RTX 30-series / RTX A-series | `sm_86` |
84
+ | L4 / L40 / L40S / RTX 40-series / RTX Ada | `sm_89` |
85
+ | GH200 / H100 / H200 | `sm_90` |
86
+ | B200 / GB200 | `sm_100` |
87
+ | RTX PRO Blackwell / GeForce RTX 50-series | `sm_120` |
88
+
89
+
90
+ ## Hello World
91
+
92
+ ```python
93
+ import numpy as np
94
+ from mojogp import RBF, SingleOutputGP
95
+
96
+ rng = np.random.default_rng(0)
97
+ X = np.linspace(-3, 3, 2000, dtype=np.float32).reshape(-1, 1)
98
+ y = (np.sin(2.0 * X[:, 0]) + 0.05 * rng.standard_normal(len(X))).astype(np.float32)
99
+
100
+ gp = SingleOutputGP(RBF())
101
+ gp.fit(X, y, max_iterations=50, method="matrix_free")
102
+
103
+ X_test = np.linspace(-4, 4, 128, dtype=np.float32).reshape(-1, 1)
104
+ mean, std = gp.predict(X_test, return_std=True, variance_method="love")
105
+ ```
106
+
107
+ ## References
108
+ 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.
109
+
110
+ 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].
111
+
112
+ 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].
113
+
114
+ 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.
115
+
116
+ 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].
117
+
118
+ 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.
119
+
120
+ 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.
121
+
122
+ 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.
123
+
124
+ 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].
125
+
126
+ 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.
127
+
128
+ 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].
129
+
130
+ 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].
131
+
132
+ Rasmussen, C.E. and Williams, C.K.I. (2008). Gaussian processes for machine learning. Cambridge, Mass. Mit Press.
133
+
134
+ 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].
135
+
136
+ 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.
137
+
138
+ 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.
139
+
140
+ 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].
141
+
142
+ 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].
143
+
144
+ 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.
145
+
146
+ ## License
147
+
148
+ MIT
@@ -0,0 +1,162 @@
1
+ """MojoGP public Python API.
2
+
3
+ The supported wrapper surface is centered on:
4
+
5
+ 1. `SingleOutputGP`
6
+ 2. `MultiOutputGP`
7
+ 3. `MultiOutputLMCGP`
8
+
9
+ Quick example:
10
+ >>> import numpy as np
11
+ >>> from mojogp import SingleOutputGP as SoGP, RBF
12
+ >>> X = np.random.randn(5000, 2).astype(np.float32)
13
+ >>> y = (np.sin(X[:, 0]) + 0.1 * np.random.randn(5000)).astype(np.float32)
14
+ >>> gp = SoGP(RBF(ard=True))
15
+ >>> gp.fit(X, y, max_iterations=50, method="matrix_free")
16
+ >>> X_test = np.random.randn(100, 2).astype(np.float32)
17
+ >>> mean, std = gp.predict(X_test, return_std=True)
18
+
19
+ Composite kernel example:
20
+ >>> from mojogp import SingleOutputGP as SoGP, RBF, Matern52
21
+ >>> kernel = RBF(active_dims=[0, 1]) + Matern52(active_dims=[1, 2])
22
+ >>> gp = SoGP(kernel)
23
+ >>> gp.fit(X, y, max_iterations=50, method="materialized")
24
+ >>> pred = gp.predict(X_test, variance_method="exact")
25
+
26
+ Multi-output example:
27
+ >>> from mojogp import MultiOutputGP, MultiOutputLMCGP
28
+ >>> Y = np.random.randn(5000, 2).astype(np.float32)
29
+ >>> icm = MultiOutputGP(kernel="rbf")
30
+ >>> icm.fit(X, Y, max_iterations=30, method="matrix_free")
31
+ >>> mean, var = icm.predict(X_test, return_var=True)
32
+ >>> lmc = MultiOutputLMCGP(kernels=["rbf", "matern52"])
33
+ >>> lmc.fit(X, Y, max_iterations=30, method="materialized")
34
+
35
+ Support boundaries worth knowing:
36
+
37
+ 1. pure categorical models are unsupported
38
+ 2. public `cholesky` posterior sampling is not part of the live API
39
+ 3. pathwise sampling is narrower than the full kernel surface
40
+ 4. `MultiOutputLMCGP(ard=True)` is available on the documented continuous LMC scope; mixed ARD applies only to continuous dimensions after categorical splitting
41
+ 5. fixed per-sample-per-task LMC observation noise `[n, T]` is available on the documented continuous LMC scope; mixed fixed noise is not public yet
42
+
43
+ The generated feature/status matrix lives in `docs/FEATURE_MATRIX.md` and is
44
+ backed by `mojogp.feature_support`.
45
+ """
46
+
47
+ # Primary API - single-output GP with JIT-compiled composite kernels
48
+ from mojogp.gp import (
49
+ SingleOutputGP,
50
+ TrainingResult,
51
+ PredictionResult,
52
+ fit_gp,
53
+ )
54
+ from mojogp._version import __version__
55
+ from mojogp.feature_support import (
56
+ ExperimentalFeatureWarning,
57
+ FeatureStatus,
58
+ InDevelopmentFeatureWarning,
59
+ MojoGPFeatureWarning,
60
+ )
61
+ from mojogp import settings as settings
62
+ from mojogp.settings import (
63
+ feature_warnings_suppressed,
64
+ get_feature_warnings_enabled,
65
+ get_progress_enabled,
66
+ progress_enabled,
67
+ set_feature_warnings_enabled,
68
+ set_progress_enabled,
69
+ )
70
+ from mojogp.progress import ProgressEvent, ProgressReporter
71
+
72
+ # Kernel specification
73
+ from mojogp.kernel import (
74
+ Kernel,
75
+ KernelNode,
76
+ KernelType,
77
+ make_ard_kernel,
78
+ compute_dim_permutation,
79
+ analyze_kernel_tree,
80
+ KernelTreeAnalysis,
81
+ CategoricalSpec,
82
+ )
83
+ from mojogp.kernel import (
84
+ RBF,
85
+ Matern12,
86
+ Matern32,
87
+ Matern52,
88
+ Periodic,
89
+ Linear,
90
+ RQ,
91
+ Polynomial,
92
+ # Categorical kernel aliases
93
+ GD,
94
+ CR,
95
+ EHH,
96
+ HH,
97
+ FE,
98
+ )
99
+
100
+ # Utilities
101
+ from mojogp.utils import StandardScaler
102
+
103
+ # Multi-output GP API
104
+ from mojogp.multi_output_gp import (
105
+ MultiOutputGP,
106
+ MultiOutputLMCGP,
107
+ MultiOutputTrainingResult,
108
+ MultiOutputPredictionResult,
109
+ LMCTrainingResult,
110
+ )
111
+
112
+ __all__ = [
113
+ # Primary API
114
+ "SingleOutputGP",
115
+ "TrainingResult",
116
+ "PredictionResult",
117
+ "fit_gp",
118
+ # Multi-output GP
119
+ "MultiOutputGP",
120
+ "MultiOutputLMCGP",
121
+ "MultiOutputTrainingResult",
122
+ "MultiOutputPredictionResult",
123
+ "LMCTrainingResult",
124
+ # Kernel specification
125
+ "Kernel",
126
+ "KernelNode",
127
+ "KernelType",
128
+ "make_ard_kernel",
129
+ # Continuous kernel shortcuts
130
+ "RBF",
131
+ "Matern12",
132
+ "Matern32",
133
+ "Matern52",
134
+ "Periodic",
135
+ "Linear",
136
+ "RQ",
137
+ "Polynomial",
138
+ # Categorical kernel shortcuts
139
+ "GD",
140
+ "CR",
141
+ "EHH",
142
+ "HH",
143
+ "FE",
144
+ # Utilities
145
+ "StandardScaler",
146
+ # Version
147
+ "__version__",
148
+ # Feature support metadata/warnings
149
+ "FeatureStatus",
150
+ "MojoGPFeatureWarning",
151
+ "ExperimentalFeatureWarning",
152
+ "InDevelopmentFeatureWarning",
153
+ "settings",
154
+ "feature_warnings_suppressed",
155
+ "get_feature_warnings_enabled",
156
+ "get_progress_enabled",
157
+ "ProgressEvent",
158
+ "ProgressReporter",
159
+ "progress_enabled",
160
+ "set_feature_warnings_enabled",
161
+ "set_progress_enabled",
162
+ ]