Gaussian processes kept close to the maths, on JAX

Low-level on purpose, and the reason is stated

The design position is in the opening paragraph and it is a claim about readability rather than about features.

GPJax aims to provide a low-level interface to Gaussian process models in JAX, structured to give researchers maximum flexibility in extending the code to suit their own needs. The stated ideal is that the code should be as close as possible to the maths written on paper when working with GP models.

That is a specific design commitment, and it explains most of what follows. A low-level library means you write the mean function, the kernel and the likelihood rather than selecting from a menu, which is exactly what a researcher wants when the method under test is a modification to one of those three things. It also means there is no convenience layer to hide behind, so the cost is on the reader.

The package is MIT licensed, authored by Thomas Pinder, and published as a paper in the Journal of Open Source Software, so it has a citable identifier as well as a version number.

The example catalogue is the evidence for the claim. It runs from an introduction for people new to the field through regression, classification and count-data regression, into sparse and stochastic variational inference, dual parameterisation, state-space formulations and multi-output models with a structured output framework, then into applied work on barycentres, graph kernels, heteroscedastic inference, orthogonal additive kernels, vector fields and ocean currents, composable spatial models, gridded data, and a benchmark dataset. Customisation guides cover building kernels, building likelihoods, deep kernels, joint inference with another probabilistic programming library, and the design of the backend module itself.

Every example exists twice, as a notebook and as a script

The examples directory holds notebooks in the double percent format, and the repository insists on it.

The stated mechanism is jupytext, so the canonical form on disk is a Python file with percent-formatted cells, and the notebook is generated from it rather than authored separately. The conversion commands are given both ways:

jupytext --to notebook example.py
jupytext --to py:percent example.ipynb

So a notebook you download from the documentation site can be turned back into a script, and a script you edit can be turned back into a notebook.

That is a small decision with real consequences. With one source of truth, a fix to an example is a pull request against a Python file that the documentation build will pick up, and there is no merge conflict between a notebook and a script describing the same computation. It also means the examples are diffable and reviewable, which notebooks famously are not.

The cost is that you cannot read the examples without either the documentation build or jupytext installed, and the percent format is not something every Python developer has seen. The documentation stack lists jupytext as a dependency for exactly this reason.

For a library whose claim is that the code should look like the maths, keeping the examples in plain Python rather than in a notebook editor's dialect is consistent.

JAX has a ceiling, and NumPy 2 is not optional

The dependency list is where a low-level library declares its real constraints, and the style is inconsistent enough to be worth reading closely.

JAX and jaxlib are pinned with a range that has both a floor and a ceiling, starting at 0.5.0 and stopping before 0.11. Everything else is a lower bound, some written with a strict greater-than and some with a floor: the gradient library above 0.2.1, the probabilistic programming library at 0.19.0, the type annotation library above 0.2.10, the progress bar above 4.66.2, runtime type checking above 0.16.1, Equinox at 0.11.0, parameterised modules from 0.0.5, a linear algebra library from 0.0.6, SciPy from 1.14, NumPy from 2.0.0, and rich from 13.0.0.

Two of those deserve comment. The JAX ceiling is the important one: a library that pins below a version is telling you it has not been tested against what comes next, and a future JAX release will fail to install rather than silently misbehave, which is the better failure. NumPy 2 rules out an environment that has not migrated to the new array API, which is a stricter floor than most projects still carry.

The rest of the list maps directly onto the library's stated purpose. Equinox is the module system the description names. A parameter handling library and a linear algebra library are what you need once you are building custom model components. A probabilistic programming library is what the joint inference guide integrates with. Runtime type checking is unusual in a scientific package and consistent with a low-level library where a wrong argument type becomes a silent numerical problem rather than an exception.

The package declares support for CPython and for PyPy, across Python 3.11, 3.12 and 3.13.

The documentation stack is copied from another repository

The optional documentation dependencies come with a comment explaining where they came from, and that comment is more informative than the dependency list.

The Sphinx toolchain is described as ported verbatim from another project, a repository called impulso, so that both share one build idiom. The stated reason is consistency across two codebases rather than any technical necessity.

The stack itself is modern: Sphinx 8.1 or later, MyST-NB for markdown with executable notebooks, the Shibuya theme, copy-button support, a design extension, bibliography support, a code auto-link extension for cross-linking tutorials to reference pages, a sitemap extension, an IPython kernel, and jupytext at 1.16 or later, with the comment noting that percent-format notebooks are read through the custom notebook formats configuration.

The one GPJax-specific difference is called out too: this project writes inline dollar-sign mathematics in its docstrings and the other one does not, which is why the markdown maths handling matters here and not there.

So the documentation build is a shared component with a small local divergence. That is a good trade for a project that does not want to maintain a second toolchain, and it does mean a documentation build failure can originate in a dependency chosen to match another repository.

Two install paths, and one command to check the result

Installation is short, and the documentation is unusually explicit about verifying it.

The stable release comes from the Python package index with a single pip command, or from the conda-forge channel, either through pixi or through conda with the channel named explicitly. Then the recommended check is one line that imports the package and prints its version:

python -c 'import gpjax; print(gpjax.__version__)'

The development version gets its own treatment, and the tone changes. It carries a warning that the version is possibly unstable and may contain bugs. The advice is to create a virtual environment first, with an example creating one on Python 3.11.0 and activating it. The setup itself clones the repository, changes into it, and creates and syncs a virtual environment with uv. The verification step is not a version print but a test run, invoking a task runner to execute all tests.

So there are two different notions of checking, matched to the two sources. A released package is verified by confirming you got the version you expected. A working copy is verified by running the test suite, which is the right check for a checkout that may be mid-change.

The stable and development paths use different tools, which is one more reason to be deliberate about which one you install.

Nothing here pins a CUDA version, which is expected for a library rather than an application.

Benchmarks, a legacy-docs redirect and two organisations

The top level shows a project with release engineering rather than a research script.

There are benchmark configuration files, a dedicated benchmarks directory, a changelog, a citation file, an agent instruction file, a context file, a static directory for assets, and a directory whose name says exactly what it does: a redirect for documentation URLs that have moved.

The benchmark configuration and its constraints file imply a specific harness, the one used for Python performance tracking with pinned environments, so performance is treated as something to measure rather than to assume. That matters for a library competing on speed in the JAX ecosystem, where the same computation can be an order of magnitude slower for a structural reason.

The legacy documentation redirect is the detail worth a second look. It exists because the documentation has moved at least once, and it means old links still resolve. For a research library whose examples are cited in papers, that is a real service rather than a nicety.

Then the organisational oddity. The repository lives under one account, while the badges, the contribution guide, the issue links, the contributor graph and the documentation site all point at a different organisation name. So the code of conduct form, the governance roles page and the slack channel belong to that other name. Before filing an issue or citing a contribution, work out which namespace is canonical.

Version 1.0.0 was tagged on 28 September 2026, after a 0.18 release in July, so the project reached a stable number recently and the changelog is where the recent history lives.

Contribution includes triage, and maintainers have a name for their role

The contribution section is more specific than most, and it names three distinct ways to help.

The first is the ordinary one: read the contributing guide, then open an issue or a discussion. The second is triaging, and the section describes what that means concretely: reproducing bug reports, asking for vital information such as version numbers and reproduction instructions, and identifying stale issues. Starting points are the open issues.

The third is where the tone changes. There is a code of conduct, contributors are expected to abide by it, and violations can be reported through a form or by reaching out to members of the project whose role has an unusual name: they are called gardeners.

A governance page defines those roles, which means the project has written down who maintains it rather than leaving it to the commit log. Combined with the stewardship metaphor, it describes maintainers as tending something rather than owning it, and the word for people reporting a problem is the opposite one.

Broader technical conversation is routed to discussions rather than issues, with best fitting practices and questions about the mathematics named as the intended subjects. That routing is a signal about the audience: the project expects questions about theory, and asks for them to go where they can be answered in public.

The slack channel is offered for development discussion and for support more generally, and the contributors list is credited explicitly.

Editorial conclusion

It fits a researcher who already knows the Gaussian process literature and wants code that maps onto the equations rather than onto someone else's class hierarchy, since the stated goal is closeness to the maths on paper. Three things to check first. JAX and jaxlib are pinned with an upper bound, so a JAX release above that ceiling will not install without editing the constraint. NumPy 2 is required, which rules out an environment pinned to the older array API. And the repository lives under one account while its badges, issue links, documentation and contributor links all point at a different organisation, so work out which is canonical before filing anything. Installation is one pip command, and a version check is recommended afterwards.

Frequently asked questions

Official sources

1. License: MIT
2. Project website
3. README
4. Releases
5. thomaspinder/GPJax on GitHub