Numpyro support?

[Chengwei94]

Is Gpjax compatible with numpyro? I have notice that there were once notebooks for numpyro and gpjax but it is no longer there. Does this means that gpjax no longer can be used with numpyro?

[daniel-dodd]

Thanks for the question @Chengwei94.

The latest GPJax releases v0.6.x don't explicitly support NumPyro in either our code or documentation. However, given we use TensorFlow probability distributions, it may be possible to wrap these (as done here https://num.pyro.ai/en/0.7.0/_modules/numpyro/contrib/tfp/distributions.html) for NumPyro compatibility.

We are keen to see a NumPyro integration though in the near future - this is on our radar.

It would be nice to e.g.,

• move all distributions from TensorFlow probability to NumPyro, or...
• ... provide a method to return NumPyro distribution's on objects instead of TensorFlow probability.

The first of these is something we tried back last Autumn #121, but there were a few issues and dropped it to push forward with releasing new features. With the new gpjax.base.Module introduced in v0.6, this should be easier to do this time round.