gpjax 0.9.2__tar.gz → 0.9.4__tar.gz

{gpjax-0.9.2 → gpjax-0.9.4}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: gpjax
-Version: 0.9.2
+Version: 0.9.4
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -19,7 +19,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: <3.13,>=3.10
 Requires-Dist: beartype>0.16.1
 Requires-Dist: cola-ml==0.0.5
-Requires-Dist: flax>=0.8.5
+Requires-Dist: flax<0.10.0
 Requires-Dist: jax<0.4.28
 Requires-Dist: jaxlib<0.4.28
 Requires-Dist: jaxopt==0.8.2
@@ -103,23 +103,23 @@ helped to shape GPJax into the package it is today.
 
 ## Notebook examples
 
-> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
-> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
-> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/examples/classification/#laplace-approximation)
-> - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/)
-> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/examples/spatial/)
-> - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/examples/barycentres/)
-> - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/examples/deep_kernels/)
-> - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/examples/poisson/)
-> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/)
+> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/_examples/regression/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/_examples/classification/)
+> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/collapsed_vi/)
+> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/uncollapsed_vi/)
+> - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/_examples/classification/#laplace-approximation)
+> - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/_examples/graph_kernels/)
+> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/_examples/spatial/)
+> - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/_examples/barycentres/)
+> - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/_examples/deep_kernels/)
+> - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/_examples/poisson/)
+> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/)
 
 ## Guides for customisation
 >
-> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/examples/yacht/)
+> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/_examples/yacht/)
 
 ## Conversion between `.ipynb` and `.py`
 Above examples are stored in [examples](docs/examples) directory in the double
@@ -180,7 +180,7 @@ optimiser = ox.adam(learning_rate=1e-2)
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=gpx.objectives.conjugate_mll,
+    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
     optim=optimiser,
     num_iters=500,

{gpjax-0.9.2 → gpjax-0.9.4}/README.md

@@ -71,23 +71,23 @@ helped to shape GPJax into the package it is today.
 
 ## Notebook examples
 
-> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
-> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
-> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/examples/classification/#laplace-approximation)
-> - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/)
-> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/examples/spatial/)
-> - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/examples/barycentres/)
-> - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/examples/deep_kernels/)
-> - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/examples/poisson/)
-> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/)
+> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/_examples/regression/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/_examples/classification/)
+> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/collapsed_vi/)
+> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/uncollapsed_vi/)
+> - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/_examples/classification/#laplace-approximation)
+> - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/_examples/graph_kernels/)
+> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/_examples/spatial/)
+> - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/_examples/barycentres/)
+> - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/_examples/deep_kernels/)
+> - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/_examples/poisson/)
+> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/)
 
 ## Guides for customisation
 >
-> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/examples/yacht/)
+> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/_examples/yacht/)
 
 ## Conversion between `.ipynb` and `.py`
 Above examples are stored in [examples](docs/examples) directory in the double
@@ -148,7 +148,7 @@ optimiser = ox.adam(learning_rate=1e-2)
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=gpx.objectives.conjugate_mll,
+    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
     optim=optimiser,
     num_iters=500,

{gpjax-0.9.2 → gpjax-0.9.4}/docs/index.md

@@ -1,4 +1,4 @@
-# Welcome to GPJax!
+# Welcome to GPJax
 
 GPJax is a didactic Gaussian process (GP) library in JAX, supporting GPU
 acceleration and just-in-time compilation. We seek to provide a flexible
@@ -6,7 +6,6 @@ API to enable researchers to rapidly prototype and develop new ideas.
 
 ![Gaussian process posterior.](static/GP.svg)
 
-
 ## "Hello, GP!"
 
 Typing GP models is as simple as the maths we
@@ -53,7 +52,7 @@ would write on paper, as shown below.
 !!! Begin
 
     Looking for a good place to start? Then why not begin with our [regression
-    notebook](https://docs.jaxgaussianprocesses.com/examples/regression/).
+    notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/).
 
 ## Citing GPJax
 

{gpjax-0.9.2 → gpjax-0.9.4}/examples/backend.py

@@ -122,7 +122,7 @@ print(constant_param._tag)
 # For most users, you will not need to worry about this as we provide a set of default
 # bijectors that are defined for all the parameter types we support. However, see our
 # [Kernel Guide
-# Notebook](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/) to
+# Notebook](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/) to
 # see how you can define your own bijectors and parameter types.
 
 # %%
@@ -156,7 +156,7 @@ transform(_close_to_zero_state, DEFAULT_BIJECTION, inverse=True)
 # may be nested within several functions e.g., a kernel function within a GP model.
 # Fortunately, transforming several parameters is a simple operation that we here
 # demonstrate for a conjugate GP posterior (see our [Regression
-# Notebook](https://docs.jaxgaussianprocesses.com/examples/regression/) for detailed
+# Notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/) for detailed
 # explanation of this model.).
 
 # %%
@@ -239,7 +239,7 @@ print(positive_reals)
 # useful as it allows us to efficiently operate on a subset of the parameters whilst
 # leaving the others untouched. Looking forward, we hope to use this functionality in
 # our [Variational Inference
-# Approximations](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/) to
+# Approximations](https://docs.jaxgaussianprocesses.com/_examples/uncollapsed_vi/) to
 # perform more efficient updates of the variational parameters and then the model's
 # hyperparameters.
 
@@ -361,7 +361,7 @@ ax.set(xlabel="x", ylabel="m(x)")
 # In this notebook we have explored how GPJax's Flax-based backend may be easily
 # manipulated and extended. For a more applied look at this, see how we construct a
 # kernel on polar coordinates in our [Kernel
-# Guide](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
+# Guide](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
 # notebook.
 #
 # ## System configuration

{gpjax-0.9.2 → gpjax-0.9.4}/examples/barycentres.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -154,9 +154,9 @@ plt.show()
 # We'll now independently learn Gaussian process posterior distributions for each
 # dataset. We won't spend any time here discussing how GP hyperparameters are
 # optimised. For advice on achieving this, see the
-# [Regression notebook](https://docs.jaxgaussianprocesses.com/examples/regression/)
+# [Regression notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/)
 # for advice on optimisation and the
-# [Kernels notebook](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/) for
+# [Kernels notebook](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/) for
 # advice on selecting an appropriate kernel.
 
 

{gpjax-0.9.2 → gpjax-0.9.4}/examples/bayesian_optimisation.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -20,7 +20,7 @@
 #
 # In this guide we introduce the Bayesian Optimisation (BO) paradigm for
 # optimising black-box functions. We'll assume an understanding of Gaussian processes
-# (GPs), so if you're not familiar with them, check out our [GP introduction notebook](https://docs.jaxgaussianprocesses.com/examples/intro_to_gps/).
+# (GPs), so if you're not familiar with them, check out our [GP introduction notebook](https://docs.jaxgaussianprocesses.com/_examples/intro_to_gps/).
 
 # %%
 from typing import (
@@ -278,7 +278,7 @@ opt_posterior = return_optimised_posterior(D, prior, key)
 # will do this using the `sample_approx` method, which generates an approximate sample
 # from the posterior using decoupled sampling introduced in ([Wilson et al.,
 # 2020](https://proceedings.mlr.press/v119/wilson20a.html)) and discussed in our [Pathwise
-# Sampling Notebook](https://docs.jaxgaussianprocesses.com/examples/spatial/). This method
+# Sampling Notebook](https://docs.jaxgaussianprocesses.com/_examples/spatial/). This method
 # is used as it enables us to sample from the posterior in a manner which scales linearly
 # with the number of points sampled, $O(N)$, mitigating the cubic cost associated with
 # drawing exact samples from a GP posterior, $O(N^3)$. It also generates more accurate

{gpjax-0.9.2 → gpjax-0.9.4}/examples/classification.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -193,15 +193,20 @@ ax.legend()
 # $\boldsymbol{x}$, we can expand the log of this about the posterior mode
 # $\hat{\boldsymbol{f}}$ via a Taylor expansion. This gives:
 #
+# $$
 # \begin{align}
 # \log\tilde{p}(\boldsymbol{f}|\mathcal{D}) = \log\tilde{p}(\hat{\boldsymbol{f}}|\mathcal{D}) + \left[\nabla \log\tilde{p}({\boldsymbol{f}}|\mathcal{D})|_{\hat{\boldsymbol{f}}}\right]^{T} (\boldsymbol{f}-\hat{\boldsymbol{f}}) + \frac{1}{2} (\boldsymbol{f}-\hat{\boldsymbol{f}})^{T} \left[\nabla^2 \tilde{p}(\boldsymbol{y}|\boldsymbol{f})|_{\hat{\boldsymbol{f}}} \right] (\boldsymbol{f}-\hat{\boldsymbol{f}}) + \mathcal{O}(\lVert \boldsymbol{f} - \hat{\boldsymbol{f}} \rVert^3).
 # \end{align}
+# $$
 #
 # Since $\nabla \log\tilde{p}({\boldsymbol{f}}|\mathcal{D})$ is zero at the mode,
 # this suggests the following approximation
+#
+# $$
 # \begin{align}
 # \tilde{p}(\boldsymbol{f}|\mathcal{D}) \approx \log\tilde{p}(\hat{\boldsymbol{f}}|\mathcal{D}) \exp\left\{ \frac{1}{2} (\boldsymbol{f}-\hat{\boldsymbol{f}})^{T} \left[-\nabla^2 \tilde{p}(\boldsymbol{y}|\boldsymbol{f})|_{\hat{\boldsymbol{f}}} \right] (\boldsymbol{f}-\hat{\boldsymbol{f}}) \right\}
 # \end{align},
+# $$
 #
 # that we identify as a Gaussian distribution,
 # $p(\boldsymbol{f}| \mathcal{D}) \approx q(\boldsymbol{f}) := \mathcal{N}(\hat{\boldsymbol{f}}, [-\nabla^2 \tilde{p}(\boldsymbol{y}|\boldsymbol{f})|_{\hat{\boldsymbol{f}}} ]^{-1} )$.

{gpjax-0.9.2 → gpjax-0.9.4}/examples/collapsed_vi.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax_beartype
 #     language: python
@@ -131,7 +131,7 @@ q = gpx.variational_families.CollapsedVariationalGaussian(
 # %% [markdown]
 # We now train our model akin to a Gaussian process regression model via the `fit`
 # abstraction. Unlike the regression example given in the
-# [conjugate regression notebook](https://docs.jaxgaussianprocesses.com/examples/regression/),
+# [conjugate regression notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/),
 # the inducing locations that induce our variational posterior distribution are now
 # part of the model's parameters. Using a gradient-based optimiser, we can then
 # _optimise_ their location such that the evidence lower bound is maximised.

{gpjax-0.9.2 → gpjax-0.9.4}/examples/constructing_new_kernels.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -71,7 +71,7 @@ cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]
 # * White noise
 # * Linear.
 # * Polynomial.
-# * [Graph kernels](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/).
+# * [Graph kernels](https://docs.jaxgaussianprocesses.com/_examples/graph_kernels/).
 #
 # While the syntax is consistent, each kernel's type influences the
 # characteristics of the sample paths drawn. We visualise this below with 10
@@ -185,7 +185,7 @@ fig.colorbar(im3, ax=ax[3], fraction=0.05)
 # We'll demonstrate this process now for a circular kernel --- an adaption of
 # the excellent guide given in the PYMC3 documentation. We encourage curious
 # readers to visit their notebook
-# [here](https://www.pymc.io/projects/docs/en/v3/pymc-examples/examples/gaussian_processes/GP-Circular.html).
+# [here](https://www.pymc.io/projects/docs/en/v3/pymc-_examples/_examples/gaussian_processes/GP-Circular.html).
 #
 # ### Circular kernel
 #
@@ -198,9 +198,15 @@ fig.colorbar(im3, ax=ax[3], fraction=0.05)
 # kernels do not exhibit this behaviour and instead _wrap_ around the boundary
 # points to create a smooth function. Such a kernel was given in [Padonou &
 # Roustant (2015)](https://hal.inria.fr/hal-01119942v1) where any two angles
-# $\theta$ and $\theta'$ are written as $$W_c(\theta, \theta') = \left\lvert
+# $\theta$ and $\theta'$ are written as
+#
+# $$
+# \begin{align}
+# W_c(\theta, \theta') & = \left\lvert
 # \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta,
-# \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
+# \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.
+# \end{align}
+# $$
 #
 # Here the hyperparameter $\tau$ is analogous to a lengthscale for Euclidean
 # stationary kernels, controlling the correlation between pairs of observations.
@@ -266,7 +272,7 @@ class Polar(gpx.kernels.AbstractKernel):
 #
 # We proceed to fit a GP with our custom circular kernel to a random sequence of
 # points on a circle (see the
-# [Regression notebook](https://docs.jaxgaussianprocesses.com/examples/regression/)
+# [Regression notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/)
 # for further details on this process).
 
 # %%

{gpjax-0.9.2 → gpjax-0.9.4}/examples/decision_making.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -22,7 +22,7 @@
 # such problems include Bayesian optimisation (BO) and experimental design. For an
 # in-depth introduction to Bayesian optimisation itself, be sure to checkout out our
 # [Introduction to BO
-# Notebook](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/).
+# Notebook](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/).
 #
 # We'll be using BO as a case study to demonstrate how one may use the decision making
 # module to solve sequential decision making problems. The goal of the decision making
@@ -76,7 +76,7 @@ cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
 # ## The Black-Box Objective Function
 #
 # We'll be using the same problem as in the [Introduction to BO
-# Notebook](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/), but
+# Notebook](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/), but
 # rather than focussing on the mechanics of BO we'll be looking at how one may use the
 # abstractions provided by the decision making module to implement the BO loop.
 #
@@ -181,7 +181,7 @@ likelihood_builder = lambda n: gpx.likelihoods.Gaussian(
 # this for us. This class takes as input a `prior` and `likeligood_builder`, which we have
 # defined above. We tend to also optimise the hyperparameters of the GP prior when
 # "fitting" our GP, as demonstrated in the [Regression
-# notebook](https://docs.jaxgaussianprocesses.com/examples/regression/). This will be
+# notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/). This will be
 # using the GPJax `fit` method under the hood, which requires an `optimization_objective`,
 # `optimizer` and `num_optimization_iters`. Therefore, we also pass these to the
 # `PosteriorHandler` as demonstrated below:
@@ -257,7 +257,7 @@ acquisition_maximizer = ContinuousSinglePointUtilityMaximizer(
 #
 # It is worth noting that `ThompsonSampling` is not the only utility function we could use,
 # since our module also provides e.g. `ProbabilityOfImprovement`, `ExpectedImprovment`,
-# which were briefly discussed in [our previous introduction to Bayesian optimisation](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/).
+# which were briefly discussed in [our previous introduction to Bayesian optimisation](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/).
 
 
 # %% [markdown]

{gpjax-0.9.2 → gpjax-0.9.4}/examples/deep_kernels.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -141,7 +141,7 @@ class DeepKernelFunction(AbstractKernel):
 # activation functions between the layers. The first hidden layer contains 64 units,
 # while the second layer contains 32 units. Finally, we'll make the output of our
 # network a three units wide. The corresponding kernel that we define will then be of
-# [ARD form](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#active-dimensions)
+# [ARD form](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#active-dimensions)
 # to allow for different lengthscales in each dimension of the feature space.
 # Users may wish to design more intricate network structures for more complex tasks,
 # which functionality is supported well in Haiku.

{gpjax-0.9.2 → gpjax-0.9.4}/examples/graph_kernels.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -22,7 +22,7 @@
 # of a graph using a Gaussian process with a Matérn kernel presented in
 # <strong data-cite="borovitskiy2021matern"></strong>. For a general discussion of the
 # kernels supported within GPJax, see the
-# [kernels notebook](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels).
+# [kernels notebook](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels).
 
 # %%
 import random
@@ -88,7 +88,9 @@ nx.draw(
 #
 # Graph kernels use the _Laplacian matrix_ $L$ to quantify the smoothness of a signal
 # (or function) on a graph
+#
 # $$L=D-A,$$
+#
 # where $D$ is the diagonal _degree matrix_ containing each vertices' degree and $A$
 # is the _adjacency matrix_ that has an $(i,j)^{\text{th}}$ entry of 1 if $v_i, v_j$
 # are connected and 0 otherwise. [Networkx](https://networkx.org) gives us an easy
@@ -151,7 +153,7 @@ cbar = plt.colorbar(sm, ax=ax)
 # non-Euclidean, our likelihood is still Gaussian and the model is still conjugate.
 # For this reason, we simply perform gradient descent on the GP's marginal
 # log-likelihood term as in the
-# [regression notebook](https://docs.jaxgaussianprocesses.com/examples/regression/).
+# [regression notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/).
 # We do this using the BFGS optimiser.
 
 # %%

{gpjax-0.9.2 → gpjax-0.9.4}/examples/intro_to_gps.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.16.6
 #   kernelspec:
 #     display_name: gpjax
 #     language: python
@@ -17,6 +17,7 @@
 # %% [markdown]
 # # New to Gaussian Processes?
 #
+#
 # Fantastic that you're here! This notebook is designed to be a gentle
 # introduction to the mathematics of Gaussian processes (GPs). No prior
 # knowledge of Bayesian inference or GPs is assumed, and this notebook is
@@ -33,10 +34,11 @@
 # model are unknown, and our goal is to conduct inference to determine their
 # range of likely values. To achieve this, we apply Bayes' theorem
 #
+# $$
 # \begin{align}
-#     \label{eq:BayesTheorem}
-#     p(\theta\,|\, \mathbf{y}) = \frac{p(\theta)p(\mathbf{y}\,|\,\theta)}{p(\mathbf{y})} = \frac{p(\theta)p(\mathbf{y}\,|\,\theta)}{\int_{\theta}p(\mathbf{y}, \theta)\mathrm{d}\theta}\,,
+#     p(\theta\mid\mathbf{y}) = \frac{p(\theta)p(\mathbf{y}\mid\theta)}{p(\mathbf{y})} = \frac{p(\theta)p(\mathbf{y}\mid\theta)}{\int_{\theta}p(\mathbf{y}, \theta)\mathrm{d}\theta},
 # \end{align}
+# $$
 #
 # where $p(\mathbf{y}\,|\,\theta)$ denotes the _likelihood_, or model, and
 # quantifies how likely the observed dataset $\mathbf{y}$ is, given the
@@ -58,7 +60,7 @@
 # family, then there exists a conjugate prior. However, the conjugate prior may
 # not have a form that precisely reflects the practitioner's belief surrounding
 # the parameter. For this reason, conjugate models seldom appear; one exception
-# to this is GP regression that we present fully in our [Regression notebook](https://docs.jaxgaussianprocesses.com/examples/regression/).
+# to this is GP regression that we present fully in our [Regression notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/).
 #
 # For models that do not contain a conjugate prior, the marginal log-likelihood
 # must be calculated to normalise the posterior distribution and ensure it
@@ -74,9 +76,13 @@
 # new points $\mathbf{y}^{\star}$ through the _posterior predictive
 # distribution_. This is achieved by integrating out the parameter set $\theta$
 # from our posterior distribution through
+#
+# $$
 # \begin{align}
 #     p(\mathbf{y}^{\star}\mid \mathbf{y}) = \int p(\mathbf{y}^{\star} \,|\, \theta, \mathbf{y} ) p(\theta\,|\, \mathbf{y})\mathrm{d}\theta\,.
 # \end{align}
+# $$
+#
 # As with the marginal log-likelihood, evaluating this quantity requires
 # computing an integral which may not be tractable, particularly when $\theta$
 # is high-dimensional.
@@ -85,13 +91,16 @@
 # distribution, so we often compute and report moments of the posterior
 # distribution. Most commonly, we report the first moment and the centred second
 # moment
+#
 # $$
 # \begin{alignat}{2}
-#     \mu  = \mathbb{E}[\theta\,|\,\mathbf{y}]  & = \int \theta p(\theta\mid\mathbf{y})\mathrm{d}\theta\\
+#     \mu  = \mathbb{E}[\theta\,|\,\mathbf{y}]  & = \int \theta
+#     p(\theta\mid\mathbf{y})\mathrm{d}\theta \quad \\
 #     \sigma^2  = \mathbb{V}[\theta\,|\,\mathbf{y}] & = \int \left(\theta -
 #     \mathbb{E}[\theta\,|\,\mathbf{y}]\right)^2p(\theta\,|\,\mathbf{y})\mathrm{d}\theta&\,.
 # \end{alignat}
 # $$
+#
 # Through this pair of statistics, we can communicate our beliefs about the most
 # likely value of $\theta$ i.e., $\mu$, and the uncertainty $\sigma$ around the
 # expected value. However, as with the marginal log-likelihood and predictive
@@ -209,9 +218,7 @@ for a, t, d in zip([ax0, ax1, ax2], titles, dists):
     d_prob = d.prob(jnp.hstack([xx.reshape(-1, 1), yy.reshape(-1, 1)])).reshape(
         xx.shape
     )
-    cntf = a.contourf(xx, yy, jnp.exp(d_prob), levels=20, antialiased=True, cmap=cmap)
-    for c in cntf.collections:
-        c.set_edgecolor("face")
+    cntf = a.contourf(xx, yy, jnp.exp(d_prob), levels=20, antialiased=True, cmap=cmap, edgecolor="face")
     a.set_xlim(-2.75, 2.75)
     a.set_ylim(-2.75, 2.75)
     samples = d.sample(seed=key, sample_shape=(5000,))
@@ -228,13 +235,16 @@ for a, t, d in zip([ax0, ax1, ax2], titles, dists):
 # %% [markdown]
 # Extending the intuition given for the moments of a univariate Gaussian random
 # variables, we can obtain the mean and covariance by
+#
 # $$
 # \begin{align}
-#    \mathbb{E}[\mathbf{y}] = \mathbf{\mu}, \quad \operatorname{Cov}(\mathbf{y}) & = \mathbf{E}\left[(\mathbf{y} - \mathbf{\mu)}(\mathbf{y} - \mathbf{\mu)}^{\top} \right]\\
+#     \mathbb{E}[\mathbf{y}] & = \mathbf{\mu}, \\
+#     \operatorname{Cov}(\mathbf{y}) & = \mathbf{E}\left[(\mathbf{y} - \mathbf{\mu})(\mathbf{y} - \mathbf{\mu})^{\top} \right] \\
 #       & =\mathbb{E}[\mathbf{y}\mathbf{y}^{\top}] - \mathbb{E}[\mathbf{y}]\mathbb{E}[\mathbf{y}]^{\top} \\
 #       & =\mathbf{\Sigma}\,.
 # \end{align}
 # $$
+#
 # The covariance matrix is a symmetric positive definite matrix that generalises
 # the notion of variance to multiple dimensions. The matrix's diagonal entries
 # contain the variance of each element, whilst the off-diagonal entries quantify
@@ -336,6 +346,7 @@ with warnings.catch_warnings():
 # $\mathbf{x}\sim\mathcal{N}(\boldsymbol{\mu}_{\mathbf{x}}, \boldsymbol{\Sigma}_{\mathbf{xx}})$ and
 # $\mathbf{y}\sim\mathcal{N}(\boldsymbol{\mu}_{\mathbf{y}}, \boldsymbol{\Sigma}_{\mathbf{yy}})$.
 # We define the joint distribution as
+#
 # $$
 # \begin{align}
 #     p\left(\begin{bmatrix}
@@ -348,6 +359,7 @@ with warnings.catch_warnings():
 #     \end{bmatrix} \right)\,,
 # \end{align}
 # $$
+#
 # where $\boldsymbol{\Sigma}_{\mathbf{x}\mathbf{y}}$ is the cross-covariance
 # matrix of $\mathbf{x}$ and $\mathbf{y}$.
 #
@@ -363,6 +375,7 @@ with warnings.catch_warnings():
 #
 # For a joint Gaussian random variable, the marginalisation of $\mathbf{x}$ or
 # $\mathbf{y}$ is given by
+#
 # $$
 # \begin{alignat}{3}
 #     & \int p(\mathbf{x}, \mathbf{y})\mathrm{d}\mathbf{y} && = p(\mathbf{x})
@@ -372,7 +385,9 @@ with warnings.catch_warnings():
 #     \boldsymbol{\Sigma}_{\mathbf{yy}})\,.
 # \end{alignat}
 # $$
+#
 # The conditional distributions are given by
+#
 # $$
 # \begin{align}
 #     p(\mathbf{y}\,|\, \mathbf{x}) & = \mathcal{N}\left(\boldsymbol{\mu}_{\mathbf{y}} + \boldsymbol{\Sigma}_{\mathbf{yx}}\boldsymbol{\Sigma}_{\mathbf{xx}}^{-1}(\mathbf{x}-\boldsymbol{\mu}_{\mathbf{x}}), \boldsymbol{\Sigma}_{\mathbf{yy}}-\boldsymbol{\Sigma}_{\mathbf{yx}}\boldsymbol{\Sigma}_{\mathbf{xx}}^{-1}\boldsymbol{\Sigma}_{\mathbf{xy}}\right)\,.
@@ -401,6 +416,7 @@ with warnings.catch_warnings():
 # We aim to capture the relationship between $\mathbf{X}$ and $\mathbf{y}$ using
 # a model $f$ with which we may make predictions at an unseen set of test points
 # $\mathbf{X}^{\star}\subset\mathcal{X}$. We formalise this by
+#
 # $$
 # \begin{align}
 #     y = f(\mathbf{X}) + \varepsilon\,,
@@ -430,6 +446,7 @@ with warnings.catch_warnings():
 # convenience in the remainder of this article.
 #
 # We define a joint GP prior over the latent function
+#
 # $$
 # \begin{align}
 #     p(\mathbf{f}, \mathbf{f}^{\star}) = \mathcal{N}\left(\mathbf{0}, \begin{bmatrix}
@@ -437,14 +454,17 @@ with warnings.catch_warnings():
 #     \end{bmatrix}\right)\,,
 # \end{align}
 # $$
+#
 # where $\mathbf{f}^{\star} = f(\mathbf{X}^{\star})$. Conditional on the GP's
 # latent function $f$, we assume a factorising likelihood generates our
 # observations
+#
 # $$
 # \begin{align}
 #     p(\mathbf{y}\,|\,\mathbf{f}) = \prod_{i=1}^n p(y_i\,|\, f_i)\,.
 # \end{align}
 # $$
+#
 # Strictly speaking, the likelihood function is
 # $p(\mathbf{y}\,|\,\phi(\mathbf{f}))$ where $\phi$ is the likelihood function's
 # associated link function. Example link functions include the probit or
@@ -453,7 +473,7 @@ with warnings.catch_warnings():
 # considers Gaussian likelihood functions where the role of $\phi$ is
 # superfluous. However, this intuition will be helpful for models with a
 # non-Gaussian likelihood, such as those encountered in
-# [classification](https://docs.jaxgaussianprocesses.com/examples/classification).
+# [classification](https://docs.jaxgaussianprocesses.com/_examples/classification).
 #
 # Applying Bayes' theorem \eqref{eq:BayesTheorem} yields the joint posterior distribution over the
 # latent function
@@ -470,7 +490,7 @@ with warnings.catch_warnings():
 # function with parameters $\boldsymbol{\theta}$ that maps pairs of inputs
 # $\mathbf{X}, \mathbf{X}' \in \mathcal{X}$ onto the real line. We dedicate the
 # entirety of the [Introduction to Kernels
-# notebook](https://docs.jaxgaussianprocesses.com/examples/intro_to_kernels) to
+# notebook](https://docs.jaxgaussianprocesses.com/_examples/intro_to_kernels) to
 # exploring the different GPs each kernel can yield.
 #
 # ## Gaussian process regression
@@ -479,20 +499,25 @@ with warnings.catch_warnings():
 # $p(y_i\,|\, f_i) = \mathcal{N}(y_i\,|\, f_i, \sigma_n^2)$,
 # marginalising $\mathbf{f}$ from the joint posterior to obtain
 # the posterior predictive distribution is exact
+#
 # $$
 # \begin{align}
 #     p(\mathbf{f}^{\star}\mid \mathbf{y}) = \mathcal{N}(\mathbf{f}^{\star}\,|\,\boldsymbol{\mu}_{\,|\,\mathbf{y}}, \Sigma_{\,|\,\mathbf{y}})\,,
 # \end{align}
 # $$
+#
 # where
+#
 # $$
 # \begin{align}
 #     \mathbf{\mu}_{\mid \mathbf{y}} & = \mathbf{K}_{\star f}\left( \mathbf{K}_{ff}+\sigma^2_n\mathbf{I}_n\right)^{-1}\mathbf{y} \\
 #     \Sigma_{\,|\,\mathbf{y}} & = \mathbf{K}_{\star\star} - \mathbf{K}_{xf}\left(\mathbf{K}_{ff} + \sigma_n^2\mathbf{I}_n\right)^{-1}\mathbf{K}_{fx} \,.
 # \end{align}
 # $$
+#
 # Further, the log of the  marginal likelihood of the GP can
 # be analytically expressed as
+#
 # $$
 # \begin{align}
 #         & = 0.5\left(-\underbrace{\mathbf{y}^{\top}\left(\mathbf{K}_{ff} + \sigma_n^2\mathbf{I}_n \right)^{-1}\mathbf{y}}_{\text{Data fit}} -\underbrace{\log\lvert \mathbf{K}_{ff} + \sigma^2_n\rvert}_{\text{Complexity}} -\underbrace{n\log 2\pi}_{\text{Constant}} \right)\,.
@@ -505,6 +530,7 @@ with warnings.catch_warnings():
 # we call these terms the model hyperparameters
 # $\boldsymbol{\xi} = \{\boldsymbol{\theta},\sigma_n^2\}$
 # from which the maximum likelihood estimate is given by
+#
 # $$
 # \begin{align*}
 #     \boldsymbol{\xi}^{\star} = \operatorname{argmax}_{\boldsymbol{\xi} \in \Xi} \log p(\mathbf{y})\,.
@@ -532,7 +558,7 @@ with warnings.catch_warnings():
 # Bayes' theorem and the definition of a Gaussian random variable. Using the
 # ideas presented in this notebook, the user should be in a position to dive
 # into our [Regression
-# notebook](https://docs.jaxgaussianprocesses.com/examples/regression/) and
+# notebook](https://docs.jaxgaussianprocesses.com/_examples/regression/) and
 # start getting their hands on some code. For those looking to learn more about
 # the underling theory of GPs, an excellent starting point is the [Gaussian
 # Processes for Machine Learning](http://gaussianprocess.org/gpml/) textbook.