Update handling of typed PRNG keys

Commit

2 years ago

Update handling of typed PRNG keys This follows the recommendations of [JEP 9263](https://github.com/google/jax/pull/17297)