Ability to return just diag on predict call (#567)

* updated linting

* linting

* reverted computations base

* fixed docstrings

* ran poe format

* changed Tuple to tuple

* changed return_cov_type to return_covariance_type

* rescoped prior lax cond and renamed functions

* revert typing Literal as beartype doesnt support it

* ran poe format

* fixed name of return function in nonconjugate posterior

* reformatted docstring args

* fixed linting error

* refactoring

* refactored and added tests

* Fix graph kernel shapes

* Fix graph kernel shapes

* Tidy typing

* Linting

* using the dense and diagonal attributes in the regression example

---------

Co-authored-by: Thomas Pinder <[email protected]>

Author: mathDR

examples/regression.py

@@ -7,7 +7,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.11.2
+#       jupytext_version: 1.17.3
 #   kernelspec:
 #     display_name: .venv
 #     language: python
@@ -121,10 +121,15 @@
 # we have just defined can be represented by a
 # [TensorFlow Probability](https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax)
 # multivariate Gaussian distribution. Such functionality enables trivial sampling, and
-# the evaluation of the GP's mean and covariance .
+# the evaluation of the GP's mean and covariance.
+#
+# Since we want to sample from the full posterior, we need to calculate the full covariance matrix.
+# We can enforce this by including the `return_covariance_type = "dense"` attribute when predicting.
+# Note this is what will be defaulted if left blank.
 
 # %%
-prior_dist = prior.predict(xtest)
+# %% [markdown]
+prior_dist = prior.predict(xtest, return_covariance_type="dense")
 
 prior_mean = prior_dist.mean
 prior_std = prior_dist.variance
@@ -212,9 +217,13 @@
 # this, we use our defined `posterior` and `likelihood` at our test inputs to obtain
 # the predictive distribution as a `Distrax` multivariate Gaussian upon which `mean`
 # and `stddev` can be used to extract the predictive mean and standard deviatation.
+# 
+# We are only concerned here about the variance between the test points and themselves, so
+# we can just copute the diagonal version of the covariance.  We enforce this by using
+# `return_covariance_type = "diagonal"` in the `predict` call.
 
 # %%
-latent_dist = opt_posterior.predict(xtest, train_data=D)
+latent_dist = opt_posterior.predict(xtest, train_data=D, return_covariance_type="diagonal")
 predictive_dist = opt_posterior.likelihood(latent_dist)
 
 predictive_mean = predictive_dist.mean