Support multidimensional HSGP (numpyro.contrib.hsgp)

[brendancooley]

• Implements eigenindices function, following Eq 10 in Riutort-Mayol et al
• Update upstream functionality to accept a $D \times m^\star$ array of eigenvalues
• Update eigenfunctions to handle multidimensional x. Vector-valued inputs are treated as unidimensional problems. Otherwise the trailing dimension of x is inferred as the dimension of the input space. All preceding dims are treated as batch dims.
• Update laplacian and approximation functions to accept list-valued inputs for m and ell. If an int (m) or float (ell) is passed when the problem is multidimensional, the same value is used for all dimensions.
• Miscellaneous cleanup and extensions of docstrings and methodology overview in "Contributed Code"

The results in @juanitorduz notebook replicate with one change to function arguments.

TODO (on this PR or future)

• [ ] support vector-valued kernel hyperparameters (alpha, length, etc)
• [ ] Update notebook in #1801 to import utility functions from numpyro.contrib.hsgp
• [ ] Update @omarfsosa birthday problem example to import utility functions from numpyro.contrib.hsgp

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[brendancooley]

^ fix pushed for python 3.9 typing issues causing CI failure

[brendancooley]

^once more, fix for linting

[juanitorduz]

This looks super cool @brendancooley ! I this one ready for review?

[juanitorduz]

Regarding the TODOS:

• The first one we can leave for other PRs so that we work on iterations (would you mind creating an issue?).
• The second would be good to change after we merge this one.
• I do not know if the third one is necessary as maybe this could serve as a "raw" example. No strong opinion. Still, we can open an issue but definitively not part of this PR.

[juanitorduz]

@fehiepsi any hints on the examples CI failing because other tests not relevant to this PR 😄 ?

[brendancooley]

@brendancooley This is amazing! I could have never done this in a cleaner way 🙌 ! Thank you!

I left two suggestions regarding documentation and some additional simple unit-tests to make this bullet-proof and to help other devs get what are the auxiliary functions expected to output. Besides that is a ✅ from my side.

Thanks the feedback @juanitorduz! Will push a commit including responses to your suggestions shortly.

[brendancooley]

@fehiepsi any hints on the examples CI failing because other tests not relevant to this PR 😄 ?

It looks like the stein_bnn example is sporadically failing to pull the boston housing data. I've been able to replicate locally but not consistently.

[brendancooley]

Regarding the TODOS:

• The first one we can leave for other PRs so that we work on iterations (would you mind creating an issue?).
• The second would be good to change after we merge this one.
• I do not know if the third one is necessary as maybe this could serve as a "raw" example. No strong opinion. Still, we can open an issue but definitively not part of this PR.

issue for vector-valued lengthscale up! https://github.com/pyro-ppl/numpyro/issues/1805

[brendancooley]

This looks great! very nice to review :D

Thanks again!