Proper splitting of PRNG key

[howsiyu]

No description provided.

[junpenglao]

Thanks!
Let's follow the official jax style and do:

key, *multiple_subkeys = random.split(key, num=k)

where key will get pass down to the next split while multiple_subkeys would be used in random call.

[howsiyu]

Changed.

[junpenglao]

There are some flaky tests - try playing around with the seed to make them past for now and leave a TODO there

[howsiyu]

I believe the current test is just lucky to pass. In fact, I have run pytest tests/mcmc/test_sampling.py::UnivariateNormalTest in the main branch with seeds from 0 to 9 and it failed 8 out of 10 times. If we really want to merge this PR we can also just change the seed to 12 and pass the test.

[codecov]

Codecov Report

Merging #526 (a50c8ec) into main (95a92fd) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #526   +/-   ##
=======================================
  Coverage   99.28%   99.28%           
=======================================
  Files          47       47           
  Lines        1948     1948           
=======================================
  Hits         1934     1934           
  Misses         14       14           

Impacted Files	Coverage Δ
blackjax/adaptation/step_size.py	100.00% <100.00%> (ø)

[junpenglao]

Thank you!