Skip to content

Commit

Permalink
[random] add more information to KeyArray deprecation error
Browse files Browse the repository at this point in the history
  • Loading branch information
jakevdp committed Oct 31, 2023
1 parent 49fedb1 commit a4e6b4e
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions jax/random.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,12 +199,16 @@
"PRNGKeyArray": (
"jax.random.PRNGKeyArray is deprecated. Use jax.Array for annotations, and "
"jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key) for runtime detection of "
"typed prng keys.", _PRNGKeyArray
"typed prng keys (i.e. keys created with jax.random.key).\n"
"For more information, see https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html",
_PRNGKeyArray
),
"KeyArray": (
"jax.random.KeyArray is deprecated. Use jax.Array for annotations, and "
"jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key) for runtime detection of "
"typed prng keys.", _PRNGKeyArray
"typed prng keys (i.e. keys created with jax.random.key).\n"
"For more information, see https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html",
_PRNGKeyArray
),
# Added September 21, 2023
"threefry2x32_key": (
Expand Down

0 comments on commit a4e6b4e

Please sign in to comment.