GPJax: a low-level Gaussian process library built on JAX and Equinox
Gaussian processes in JAX and Equinox.
At a glance
- What is it?
- GPJax exposes GP modelling primitives in JAX so that researchers can write kernels, likelihoods and inference loops close to the mathematics on paper. The trade-off is that it hands you the pieces, not a finished pipeline.
- Who is it for?
- Adopt GPJax if you already work in JAX and need to modify the GP itself, not just fit one: custom kernels, non-Euclidean inputs, graph kernels, heteroscedastic noise, or a variational objective you want to write yourself. Do not adopt it if you need a batteries-included regression API with plotting and diagnostics, or if your team has no JAX experience, because the README's own framing is a low-level interface aimed at researchers.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 2 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap GPJax is trying to fill
Most Gaussian process libraries optimise for the common case: give them a design matrix and a target vector, and they return a fitted model. That is the right shape for applied work, but it is the wrong shape when the model itself is the research object. If you want a kernel that is not in the catalogue, a likelihood that is not Gaussian, or an inference scheme that does not match the library's default, you end up either subclassing something deeply nested or reimplementing the GP from scratch. GPJax takes the second option and makes it the design. The README states the aim directly: to provide a low-level interface to Gaussian process models in JAX, structured to give researchers maximum flexibility in extending the code, with the code kept as close as possible to the mathematics written on paper. The intended audience is therefore people who are comfortable with the algebra of GPs and want the library to stay out of the way. It is not aimed at someone who wants a fitted curve in five lines, although the quickstart example suggests that is possible for simple conjugate regression.
What sits underneath: JAX, Equinox and a functional core
The two dependencies named in the repository description do most of the architectural work. JAX supplies automatic differentiation, vectorisation and just-in-time compilation, which means a GP model here is a Python function that JAX can transform rather than a mutable object with a fit method. Equinox supplies the module and parameter container layer on top, which is why the library can describe a model as a PyTree and still let you write it with familiar class syntax. The practical consequence is that gradients and compilation come from JAX, so anything you write has to be JAX-traceable: no in-place mutation of arrays, no data-dependent Python control flow inside a compiled region, and shapes fixed at trace time. That constraint is the price of the flexibility. The repository layout reinforces the low-level framing. The examples live in docs/examples as py:percent files, with jupytext used to convert between .py and .ipynb, which suggests the notebooks are generated from scripts rather than maintained as the primary artefact. The README's package organisation section is short and points outward to a contributing guide, a code of conduct and a governance page listing roles such as gardeners, so the project has an explicit decision-making structure rather than a single maintainer acting alone.
Getting it installed, stable or from source
The stable path is a single command: pip install gpjax. Conda users have two options listed, pixi add gpjax or conda install --channel conda-forge gpjax, both drawing from the conda-forge feedstock. The README recommends checking the installed version with python -c 'import gpjax; print(gpjax.__version__)', which is worth doing because the release history shows a fast cadence: v0.15.0 in June 2026, v0.17.0 in July, v0.18.0 later that same month. For the development version the README warns explicitly that it is possibly unstable and may contain bugs, advises creating a virtual environment first (the example uses conda create -n gpjax_experimental python=3.11.0), then gives the sequence git clone, cd GPJax, uv venv, uv sync. It also recommends running the supplied unit tests with uv run poe all-tests, which implies the project uses poe as a task runner and uv for environment management. Note the two ecosystems coexisting here: conda for the documented virtual environment, uv for the actual dependency sync. Nothing in the material explains why, and it is a small inconsistency a new contributor will notice.
The inference settings the documentation actually covers
The README lists eleven notebook examples, and the list is the most useful signal about scope. Conjugate inference and classification cover the standard cases. Sparse variational inference and stochastic variational inference are separate notebooks, which matters because they are different algorithms with different scaling behaviour and different failure modes; a library that documents them separately is telling you to choose deliberately. The Laplace approximation is documented as a section within the classification notebook rather than as its own page, so it is a lighter-weight offering than the variational routes. Beyond that the list moves into territory where most GP libraries stop: inference on non-Euclidean spaces, inference on graphs, heteroscedastic inference, learning Gaussian process barycentres, deep kernel regression and Poisson regression. Custom kernels get their own guide, with a custom-kernel anchor inside the constructing-new-kernels notebook, and there is a UCI regression walkthrough using the yacht dataset. The breadth is real, but breadth in a notebook list is not the same as depth in an API reference. Whether each of these is a maintained code path or a demonstration is something you can only settle by reading the notebooks themselves.
Where the low-level choice costs you
The design decision that makes GPJax extensible is the same one that makes it demanding. There is no evidence in the supplied material of a high-level estimator interface, a model-selection helper, or built-in diagnostic and plotting utilities. The README does not claim any of these, so their absence should not be read as a failing, but it does define the boundary. If your task is fitting a GP to a few thousand rows with a squared exponential kernel and reporting predictive intervals, you will write more code here than you would elsewhere, and that code will be your responsibility to test. The second cost is the JAX constraint. Debugging a traced function is harder than debugging ordinary Python, and error messages from inside a jit-compiled region are frequently about shapes rather than about the modelling mistake you actually made. The third is version churn. Three minor releases in roughly three months, with the development branch flagged as possibly unstable, means pinning your dependency is not optional. A pinned version plus the version-check command from the README is the minimum discipline.
GPyTorch and the difference in approach
The obvious comparison is GPyTorch, which builds on PyTorch and organises GP modelling around a module hierarchy with an explicit training loop and a strong emphasis on scalable approximations such as inducing-point methods and structured kernels for grid data. The difference is not which features exist but where the abstraction sits. GPyTorch gives you a model class you configure and train; the extension points are subclasses and registered components within a framework that decides how batches, caches and variational objectives fit together. GPJax gives you functions and PyTrees that JAX transforms, so the composition is yours to define and the compiled artifact is a product of your code rather than the library's. That makes GPJax more natural when the model is unusual and less natural when it is standard. There is a secondary difference in the surrounding ecosystem: GPyTorch inherits PyTorch's tooling for debugging and profiling, while GPJax inherits JAX's compilation model, which is faster when it works and more opaque when it does not. Neither is strictly better; they fail in different places.
Licence, maintenance and what a version bump implies
The licence is MIT, which is permissive and imposes few obligations beyond retaining the copyright notice and permission text. That matters for commercial use and for forking, and it is consistent with the project's research framing. This is not legal advice; read the LICENSE file in the repository if the terms affect a product decision. On maintenance, the signals available are the push date of September 2026, the release cadence across mid-2026, and the governance page referenced from the README, which describes roles and a code of conduct. Those suggest an active project with more than one person involved, though the material does not name maintainers or describe release policy. The upgrade cost is the practical concern. With minor versions arriving roughly monthly, an upgrade can change API surface, and the README's own warning about the development version being possibly unstable applies with less force but not zero force to tagged releases. The mitigation is in the repository already: uv run poe all-tests runs the supplied test suite, so an upgrade can be validated against the project's own expectations before it reaches your model code. If you use GPJax in published work, the JOSS paper citation is given in the README as a BibTeX entry with DOI 10.21105/joss.04455.
Editorial conclusion
Adopt GPJax if you already work in JAX and need to modify the GP itself, not just fit one: custom kernels, non-Euclidean inputs, graph kernels, heteroscedastic noise, or a variational objective you want to write yourself. Do not adopt it if you need a batteries-included regression API with plotting and diagnostics, or if your team has no JAX experience, because the README's own framing is a low-level interface aimed at researchers. Before committing, verify the installed version with python -c 'import gpjax; print(gpjax.__version__)', run uv run poe all-tests against the development checkout, and read the notebook that matches your inference setting (collapsed_vi, uncollapsed_vi, heteroscedastic_inference) to confirm the API surface you need is actually documented rather than implied.
Community notes