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.

[thomaspinder]

The challenge I am having right now is working out a clean way to map the samples generated by numpyro back into a Flax/NNX state. If this were possible, then the prediction step would work.

[juanitorduz]

I am also looking into this 😄. I'll need to sleep it over 😅. Keep me posted 💪

[thomaspinder]

Likewise, should you have any productive ideas here!

[thomaspinder]

google/flax#4704 (comment)

[juanitorduz]

@thomaspinder In juanitorduz@b42ed96 I have a brute force approach to parse the samples into the state template (I am pretty sure there should be a nicer version using nnx.transversals as suggested) (Update: juanitorduz@8e1209d#commitcomment-155494638)