New jax keys
reubenharry opened this issue · 1 comments
reubenharry commented
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
junpenglao commented
We are already using the new random key: #569. It is just we still call the custom type PRNGKey