GPJax

GPJax provides a low-level interface to Gaussian process (GP) models built on Jax. It is designed for researchers, aiming to keep the code as close as possible to the mathematical formulations used in GP modeling.

Supported inference methods and use cases include:

• Inference Methods: Conjugate Inference, Sparse Variational Inference, Stochastic Variational Inference, Laplace Approximation.
• Regression/Classification: Regression, Classification, Poisson Regression, Heteroscedastic Inference, Deep Kernel Regression.
• Advanced Topics: Inference on Non-Euclidean Spaces, Inference on Graphs, Learning Gaussian Process Barycentres.

GPJax is a didactic Gaussian Process (GP) library designed for researchers to develop novel GP methodology. It provides a set of composable objects that closely resemble mathematical notation, making it easier to translate theory into code.

Built on top of JAX, GPJax natively supports:

• Efficient compilation to XLA
• Automatic differentiation
• Vectorised operations
• Hardware acceleration on CPUs, GPUs, and TPUs

Because GPJax objects are built using Chex dataclasses and registered as PyTree nodes, they are fully compatible with standard JAX transformations like jit, grad, and vmap.

The gpjax.state_space module provides tools for working with State-Space Gaussian Processes. This includes defining state-space priors, computing state-space marginal likelihoods, and performing model fitting using various optimizers.

Key components available in this module include:

• Priors: StateSpacePrior for defining the GP prior in state-space form.
• Posterior Inference: StateSpaceConjugatePosterior for efficient conjugate inference.
• Fitting: fit, fit_lbfgs, and fit_scipy for optimizing model hyperparameters.
• Utilities: state_space_mll for calculating the marginal log-likelihood and to_sde for converting models to Stochastic Differential Equations.

GPJax is designed to provide an accurate representation of Gaussian process mathematics. Variable names in the API are chosen to closely match the notation used in Rasmussen & Williams (2006).

When reading GPJax code or documentation, use the following mapping to relate code variables to mathematical quantities:

Standard Gaussian Process Notation

• n: Number of training inputs
• x: Training inputs
• y: Training labels
• t: Test inputs
• f: Latent function modelled as a GP
• fx: Latent function at inputs x
• mux: Prior mean at inputs x
• Kxx: Kernel Gram matrix at inputs x
• Lx: Lower Cholesky factor of Kxx
• Ktx: Cross-covariance between inputs t and x

Sparse Gaussian Process Notation

• m: Number of inducing points
• z: Inducing inputs
• u: Inducing outputs

GPJax supports parameterisation using natural parameters $(\mathbf{\Sigma}^{-1}\boldsymbol{\mu}, -\tfrac{1}{2}\mathbf{\Sigma}^{-1})$ instead of moment parameterisation $(\boldsymbol{\mu}, \mathbf{\Sigma})$. Using natural parameters makes natural gradients computationally cheap (avoiding Fisher-matrix inversion) and typically leads to much faster convergence than plain gradient descent on moments.

Relevant classes include:

• NaturalVariationalGaussian
• ExpectationVariationalGaussian

GPJax follows a mathematical composition pattern. A Gaussian Process posterior is represented as the product of a Prior and a Likelihood:

posterior = prior * likelihood

This abstraction maps directly to the Bayesian identity $p(f|y) ext{ } ext{'} ext{ } p(f)p(y|f)$, where the Prior represents $p(f)$ and the Likelihood represents $p(y|f)$.

GPJax is designed to integrate seamlessly with the broader JAX ecosystem for various probabilistic and machine learning tasks:

• Inference (MCMC): Can utilize samplers from BlackJax and TensorFlow Probability.
• Optimisation: Integrates with Optax for gradient-based optimisation, providing access to various optimisers and learning rate schedules.
• Probability Distributions: Leverages Distrax and TensorFlow Probability for efficient representation.
• Deep Learning: Can incorporate functionality from Haiku to combine GPs with deep learning methods.

When the marginal log-likelihood is unavailable (e.g., with non-conjugate likelihoods), GPJax uses the Evidence Lower Bound (ELBO). There are two primary implementations:

1. elbo: The uncollapsed, mini-batchable bound.
2. collapsed_elbo: The collapsed bound, which solves variational parameters analytically. This requires a conjugate likelihood and a full pass over the data.