Use public JAX API for PRNGKeyArray

[junpenglao]

Thank you for opening a PR!

A few important guidelines and requirements before we can merge your PR:

• If I add a new sampler, there is an issue discussing it already;
• We should be able to understand what the PR does from its title only;
• There is a high-level description of the changes;
• There are links to all the relevant issues, discussions and PRs;
• The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• There are tests covering the changes;
• The doc is up-to-date;
• If I add a new sampler* I added/updated related examples

Consider opening a Draft PR if your work is still in progress but you would like some feedback from other contributors.

[codecov]

Codecov Report

Merging #498 (3224185) into main (6b76746) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #498   +/-   ##
=======================================
  Coverage   99.16%   99.16%           
=======================================
  Files          48       48           
  Lines        1915     1915           
=======================================
  Hits         1899     1899           
  Misses         16       16           

Impacted Files	Coverage Δ
blackjax/adaptation/mass_matrix.py	100.00% <ø> (ø)
blackjax/adaptation/meads_adaptation.py	100.00% <ø> (ø)
blackjax/adaptation/step_size.py	100.00% <ø> (ø)
blackjax/adaptation/window_adaptation.py	100.00% <ø> (ø)
blackjax/diagnostics.py	100.00% <ø> (ø)
blackjax/kernels.py	99.17% <ø> (ø)
blackjax/mcmc/elliptical_slice.py	95.00% <ø> (ø)
blackjax/mcmc/integrators.py	100.00% <ø> (ø)
blackjax/mcmc/irmh.py	100.00% <ø> (ø)
blackjax/mcmc/metrics.py	100.00% <ø> (ø)
... and 14 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.