Breaking for jax 0.4.24
Closed this issue · 0 comments
Intron7 commented
Can you update chex to work with the newest version of jax.
From jax.random: PRNGKeyArray, KeyArray,
default_prng_impl, threefry_2x32, threefry2x32_key, threefry2x32_p, rbg_key, and unsafe_rbg_key.
This breaks the completely breaks the import of chex and makes many packages unusable.