Add docstring to pymc.sampling.jax.sample_jax_nuts

[andrewdipper]

Add docstring to pymc.sampling.jax.sample_jax_nuts

Description

As above - mostly just copied from _sample_blackjax_nuts. Fixes empty documentation here https://www.pymc.io/projects/docs/en/stable/api/generated/pymc.sampling.jax.sample_blackjax_nuts.html#pymc.sampling.jax.sample_blackjax_nuts and here https://www.pymc.io/projects/docs/en/stable/api/generated/pymc.sampling.jax.sample_numpyro_nuts.html#pymc.sampling.jax.sample_numpyro_nuts

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7313.org.readthedocs.build/en/7313/

[welcome]

]
Congrats on merging your first pull request! 🎉 We here at PyMC are proud of you! 💖 Thank you so much for your contribution 🎁

[ricardoV94]

Thanks @andrewdipper