blackjax-devs/blackjax

New jax keys

reubenharry opened this issue · 1 comments

Current behavior

We currently use jax.random.PRNGKey throughout.

Desired behavior

Following this jax discussion https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html, I'm wondering if we want to switch to the new key style, jax.random.key

We are already using the new random key: #569. It is just we still call the custom type PRNGKey