GPJax#

Gaussian processes in JAX.

GPJax is a didactic Gaussian process (GP) library in JAX, supporting GPU acceleration and just-in-time compilation. We seek to provide a flexible API to enable researchers to rapidly prototype and develop new ideas.

Gaussian process posterior.

“Hello, GP!”#

Typing GP models is as simple as the maths we would write on paper.

import gpjax as gpx

mean = gpx.mean_functions.Zero()
kernel = gpx.kernels.RBF()
prior = gpx.gps.Prior(mean_function=mean, kernel=kernel)
likelihood = gpx.likelihoods.Gaussian()

model = prior * likelihood  # the joint p(f, y)

Conditioning the model on data yields the posterior process, which can then be queried at any test inputs:

import jax.numpy as jnp

xtrain = jnp.linspace(0.0, 1.0, 20).reshape(-1, 1)
D = gpx.Dataset(X=xtrain, y=jnp.sin(xtrain))
xtest = jnp.linspace(0.0, 1.0, 50).reshape(-1, 1)

posterior = model.condition(D)  # p(f | D) — equivalently: model | D
predictive = posterior(xtest)

(1)#\[\begin{split} \begin{aligned} k(\cdot, \cdot') & = \sigma^2\exp\left(-\frac{\lVert \cdot- \cdot'\rVert_2^2}{2\ell^2}\right)\\ p(f(\cdot)) & = \mathcal{GP}(\mathbf{0}, k(\cdot, \cdot')) \\ p(y\,|\, f(\cdot)) & = \mathcal{N}(y\,|\, f(\cdot), \sigma_n^2) \\ \\ p(f(\cdot) \,|\, y) & \propto p(f(\cdot))p(y\,|\, f(\cdot))\,. \end{aligned} \end{split}\]

We currently have some availability for consulting on how Gaussian processes, Bayesian modelling, and GPJax can be integrated into your team's work. If this sounds relevant to your work, book an introductory call. These calls are for consulting inquiries only. For technical usage questions and free community support, please use GitHub Discussions and the documentation below.

Learn more#

Installation

Install the stable or development version.

Installation

New to Gaussian processes?

Priors, posteriors and the marginal likelihood from first principles.

New to Gaussian Processes?

Regression

The canonical end-to-end workflow, start to finish.

Regression

API reference

Every public class and function, with source links.

API Reference

Sharp bits

The numerical pitfalls worth knowing about before you hit them.

The sharp bits

Benchmarks

The ASV dashboard tracking GPJax’s performance commit by commit.

benchmarks/index.html

Citing GPJax#

If you use GPJax in your research, please cite our JOSS paper.

@article{Pinder2022,
  doi = {10.21105/joss.04455},
  url = {https://doi.org/10.21105/joss.04455},
  year = {2022},
  publisher = {The Open Journal},
  volume = {7},
  number = {75},
  pages = {4455},
  author = {Thomas Pinder and Daniel Dodd},
  title = {GPJax: A Gaussian Process Framework in JAX},
  journal = {Journal of Open Source Software}
}