Reduce blackjax sampling memory usage

[junpenglao]

Description

Reduce blackjax sampling memory usage by not outputting the warm up diagnostics

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--7407.org.readthedocs.build/en/7407/

[junpenglao]

Will need to upgrade jaxlib requirement first (conda-forge/jaxlib-feedstock#272)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.18%. Comparing base (641a60b) to head (e81d828).
Report is 97 commits behind head on main.

Additional details and impacted files

@@ Coverage Diff @@ ## main #7407 +/- ## ======================================= Coverage 92.18% 92.18% ======================================= Files 103 103 Lines 17258 17259 +1 ======================================= + Hits 15909 15910 +1  Misses 1349 1349 

Files with missing lines	Coverage Δ
pymc/sampling/jax.py	94.03% <100.00%> (+0.02%)	⬆️