GPJax: A Gaussian Process Framework in JAX

Abstract

Gaussian processes (GPs, Rasmussen & Williams, 2006) are Bayesian nonparametric models that have been successfully used in applications such as geostatistics 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 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.