Sample stats for blackjax nuts

[TimOliverMaier]

What is this PR about?
With this PR I want to address #5718. As of now no stats are saved in InferenceData if one uses sample_blackjax_nuts.
However I am unsure about the meaning of the attributes energy and num_trajectory_expansions.
Here is what is documented in the blackjax code. I first understood energy to be the absolute energy of a state and num_trajectory_expansions to be the tree_depth (That's the current state of the code). But second reading made me skeptical. Any help here is very welcome 😃

How it works:

• blackjax returns sample statistics as instance of blackjax.mcmc.nuts.NUTSInfo. The attributes of this instance
  are converted to a dictionary Dict[str,ArrayLike] mapping the stats with pymc naming conventions
• potential energy is stored alongside states in a blackjax.mcmc.hmc.HMCState object. The potential energy is extracted and
  stored in the above mentioned dictionary.
• the dictionary is passed as sample_stats argument to arviz.from_dict at the end of pymc.sampling.jax.sample_blackjax_nuts

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 👇

Bugfixes / New features

• pm.sample.jax.sample_blackjax_nuts returns InferenceData with sample_stats group

[codecov]

Codecov Report

Merging #6264 (5f02ad9) into main (b5db350) will increase coverage by 6.16%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6264      +/-   ##
==========================================
+ Coverage   88.15%   94.31%   +6.16%     
==========================================
  Files         111      111              
  Lines       23820    23855      +35     
==========================================
+ Hits        20998    22499    +1501     
+ Misses       2822     1356    -1466     

Impacted Files	Coverage Δ
pymc/sampling/jax.py	97.34% <100.00%> (+0.14%)	⬆️
pymc/tests/sampling/test_jax.py	100.00% <100.00%> (ø)
pymc/model.py	89.51% <0.00%> (+0.27%)	⬆️
pymc/distributions/shape_utils.py	97.85% <0.00%> (+0.42%)	⬆️
pymc/backends/arviz.py	90.08% <0.00%> (+2.89%)	⬆️
pymc/distributions/logprob.py	97.26% <0.00%> (+6.16%)	⬆️
pymc/step_methods/hmc/quadpotential.py	80.69% <0.00%> (+6.93%)	⬆️
pymc/model_graph.py	78.82% <0.00%> (+12.35%)	⬆️
pymc/variational/updates.py	92.11% <0.00%> (+54.18%)	⬆️
pymc/distributions/timeseries.py	94.45% <0.00%> (+65.99%)	⬆️
... and 4 more

[junpenglao]

Your interpretation of the mapping between the different naming is correct here. LGTM.

[TimOliverMaier]

Thank you @junpenglao . The failing test succeded locally. Will look into it asap

[TimOliverMaier]

The test failed because I had a newer version of blackjax installed, that stores n_steps (called num_integration_steps there) . The test now no longer asserts the existence of a DataArray n_steps in sample_stats. The code however would store n_steps if a user has a newer version of blackjax installed.

[junpenglao]

Looks good - you will need to rebase the change.

[ricardoV94]

@TimOliverMaier can you add (and fill) the relevant bullet point from the PR template to your top comment? We use those for the auto-generated release notes: https://github.com/pymc-devs/pymc/blob/main/.github/PULL_REQUEST_TEMPLATE.md

[TimOliverMaier]

In the last commit I only moved the sample_stats argument into the call of partial. That's where other groups are passed and that's how it is done in the numpyro nuts function.

[michaelosthege]

Looks like this is ready to merge, so let's do it!