GPJax: A Gaussian Process Framework in JAX
Thomas Pinder, Daniel J. Dodd · The Journal of Open Source Software · 2022
Gaussian processes (GPs, Rasmussen & Williams, 2006) are Bayesian nonparametric models that have been successfully used in applications such as geostatistics (Matheron, 1963), Bayesian optimisation (Mockus et al., 1978), and reinforcement learning (Deisenroth & Rasmussen, 2011).GPJax is a didactic GP library targeted at researchers who wish to develop novel GP methodology.The scope of GPJax is to provide users with a set of composable objects for constructing GP models that closely resemble the underlying maths that one would write on paper.Furthermore, by the virtue of being written in JAX (Bradbury et al., 2018), GPJax natively supports CPUs, GPUs and TPUs through efficient compilation to XLA, automatic differentiation and vectorised operations.Consequently, GPJax provides a modern GP package that can effortlessly be tailored, extended and interleaved with other libraries to meet the individual needs of researchers and scientists. Statement of NeedFrom both an applied and methodological perspective, GPs are widely employed in the statistics and machine learning communities.High-quality software packages that promote GP modelling are accountable for much of their success.However, there currently exists a gap within the JAX ecosystem for a Gaussian process package to be developed that incorporates scalable inference techniques.GPJax seeks to resolve this.