New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

chex: alias PRNGKey to jax.Array #254

Open

copybara-service wants to merge 1 commit into master from test_565133147

Commits on Sep 13, 2023

chex: alias PRNGKey to jax.Array

Going forward, `jax.Array` is the correct type annotation for both new-style and old-style PRNG keys in JAX (see jax-ml/jax#17297)

Note that `jax.random.KeyArray` has been aliased to `Any` under TYPE_CHECKING, and so this change will make existing annotations more strict, not less strict.

PiperOrigin-RevId: 565133147

Jake VanderPlas authored and DistraxDev committed Sep 13, 2023