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.