Gaussian processes that actually look like the math

What it does

GPJax is a Gaussian process modeling library built on JAX and Equinox. It provides the standard suite of GP inference methods—conjugate regression, variational approximations, classification, and Bayesian optimization—through an intentionally low-level API. The goal is to let researchers tweak kernels and inference schemes without wading through layers of framework indirection.

The interesting bit

Most GP libraries hide the matrix algebra behind high-level objects; GPJax does the opposite, keeping the algebra visible so the code reads like the LaTeX. That transparency makes it easier to prototype unusual ideas—non-Euclidean kernels, graph-structured data, or deep kernel learning—without fighting the framework’s abstractions.

Key highlights

• Conjugate, sparse variational, stochastic variational, and Laplace inference
• Non-Euclidean spaces, graph kernels, heteroscedastic noise, and deep kernel regression
• Bayesian optimization and Poisson regression examples
• Published in the Journal of Open Source Software (JOSS)
• Built on JAX and Equinox

Verdict

Worth a look for researchers who need to customize kernels or inference schemes. Skip it if you just need a one-line GP fit and do not care about the underlying algebra.