Add GPJax notebook

[thomaspinder]

GPJax is a library built upon Numpyro and Flax's NNX module that allows for flexible Gaussian processes to be implemented. Within the past 6 months, a series of PRs to GPJax have brought it closer to numpyro and it can now be easily integrated into a larger numpyro model. This PR demonstrates how people can achieve that by leveraging the synergies present between Numpyro and GPJax.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[thomaspinder]

cc @juanitorduz / @theorashid

[theorashid]

Nice, yeah we should get this in.

Immediate responses:

• Let's try and use random_nnx_module instead of register_parameters, as discussed offline
• gpjax shouldn't become a dependency. Can you make the import optional. EDIT: dev dependency is probably fine, but lower pin the version for when you integrated numpyro
• We could also convert this to one of the short py file examples? Keeps it lighter, it can link to your GPJax docs for full depth. But i'm easy either way.

[thomaspinder]

• So random_nnx_module won't, unfortunately, work as it operates on NNX's Param class whereas GPJax subclasses NNX Variable. We must use Variable in GPJax, not Param , in order to attach bijectors to a parameter. I wrote a short FAQ here earlier today.
• Good call on the pin. Will set.
• Keeping the notebook feels nice(r) for now, but I'm also open to either.

[fehiepsi]

Very cool! Thanks for contributing!

[thomaspinder]

CI / examples (3.13) appears to be failing, but I'm not sure it's related to my change @fehiepsi . Any thoughts?

[Qazalbash]

CI / examples (3.13) appears to be failing, but I'm not sure it's related to my change @fehiepsi . Any thoughts?

We are encountering it on other PRs too. I presume the dataset download URL is inaccessible. We are investigating this issue.

[fehiepsi]

Could you sync to master?

[juanitorduz]

can we add

!pip install -q numpyro@git+https://github.com/pyro-ppl/numpyro

as other notebooks? See https://num.pyro.ai/en/stable/tutorials/bayesian_regression.html

Also the line

assert numpyro.__version__.startswith("0.20.0")

Reply via ReviewNB

[juanitorduz]

Suggestion! Show the trace and show that we have recovered the parameters, for example like in https://num.pyro.ai/en/stable/tutorials/other_samplers.html

Reply via ReviewNB

[juanitorduz]

Thanks @thomaspinder ! I left two minor comments. From my side, it's good to go :)

I fixed the unrelated errors in #2151 so after you rebase, this should be fine :)

[thomaspinder]

Nice idea @juanitorduz! I have addressed both of your comments.

[juanitorduz]

Amazing! Thanks!