chex: alias PRNGKey to jax.Array#254

copybara-service[bot] wants to merge 1 commit intomaingoogle-deepmind/distrax:mainfrom

test_565133147google-deepmind/distrax:test_565133147Copy head branch name to clipboard