Improve blackjax sampling integration

[junpenglao]

Closes #6550

What is this PR about?
Adding more fine tune control in blackjax sampling.

After this PR:

# Use full rank mass matrix in window adaptation 
with model:
    idata2 = pm.sample(
        nuts_sampler="blackjax", 
        nuts_sampler_kwargs={
            "chain_method": "vectorized",
            "adaptation_kwargs": {
                "is_mass_matrix_diagonal": False
                }
            }
        )

# Warm up and sample with HMC (instead of NUTS) with full rank mass matrix adaptation
with model:
    idata2 = pm.sample(
        nuts_sampler="blackjax", 
        nuts_sampler_kwargs={
            "chain_method": "vectorized",
            "adaptation_kwargs": {
                "is_mass_matrix_diagonal": False,
                "algorithm": "hmc",
                "num_integration_steps": 20,
                }
            }
        )

# Enable progress bar (only works for `chain_method =  "vectorized"`)
with model:
    idata2 = pm.sample(
        nuts_sampler="blackjax", 
        progress_bar=True,
        nuts_sampler_kwargs={
            "chain_method": "vectorized",
            # "chain_method": "parallel",
            "adaptation_kwargs": {
                "is_mass_matrix_diagonal": False,
                }
            }
        )

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

None

New features

More option in Blackjax sampling

Bugfixes

None

Documentation

None

Maintenance

None

---

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

[codecov]

Codecov Report

Merging #6963 (826cbe8) into main (5f29b25) will decrease coverage by 0.03%.
Report is 2 commits behind head on main.
The diff coverage is 80.43%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6963      +/-   ##
==========================================
- Coverage   92.12%   92.09%   -0.03%     
==========================================
  Files         100      100              
  Lines       16859    16875      +16     
==========================================
+ Hits        15531    15541      +10     
- Misses       1328     1334       +6     

Files	Coverage Δ
pymc/sampling/mcmc.py	87.69% <ø> (ø)
pymc/sampling/jax.py	93.07% <80.43%> (-3.24%)	⬇️

... and 1 file with indirect coverage changes

[junpenglao]

@jessegrabowski FYI the progress bar is partially working for both warm up and sampling.