Warn on divergences after sampling with JAX

[jessegrabowski]

What is this PR about?
Closes #7041

This is a draft to add convergence checks to the JAX samplers.

Right now I'm just calling run_convergence_checks after sampling. It might be nice to instead wrap the JAX returns in _sample_return, but this was easier. Feedback requested.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• The blackjax sampler should now issue a warning on divergences

New features

• None

Bugfixes

• None

Documentation

• None

Maintenance

• None

---

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

[codecov]

Codecov Report

Merging #7051 (e28c71e) into main (005ba5f) will increase coverage by 0.00%.
Report is 1 commits behind head on main.
The diff coverage is 100.00%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #7051   +/-   ##
=======================================
  Coverage   92.16%   92.16%           
=======================================
  Files         101      101           
  Lines       16827    16831    +4     
=======================================
+ Hits        15509    15513    +4     
  Misses       1318     1318           

Files	Coverage Δ
pymc/sampling/mcmc.py	87.79% <100.00%> (+0.10%)	⬆️

[ricardoV94]

Looks fine, but why not in the sample_numpyro_nuts function itself? I am thinking of that because some users use that directly and also we already re-implement most of the logic of _sample_return there?

[jessegrabowski]

I whipped up this as a live example during the sprint. I'll go back and do a better job following your suggestion @ricardoV94