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).

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	Coverage Δ
pymc/sampling/jax.py	94.03% <100.00%> (+0.02%)	⬆️