Type annotation for PRNG keys

[carlosgmartin]

What's the recommended type annotation for a PRNG key (as opposed to a generic Array)?

I think it used to be PRNGKey, but that's legacy.

[jakevdp]

The right annotation is jax.Array – this was decided here: https://docs.jax.dev/en/latest/jep/9263-typed-keys.html#type-annotations-for-prng-keys

New-style typed keys (produced by jax.random.key) are defined by their dtype, and JAX does not offer any dtype-specific array annotations, so Array is the only appropriate annotation. Old-style keys (produced by jax.random.PRNGKey) are just arrays with dtype uint32 and a trailing dimension of a particular size depending on the default PRNG impl, and so Array is the only appropriate annotation.

[jakevdp]

Note also that jax.random.PRNGKey is a function, not a type, so the only people who ever used it as an annotation were folks who were not running a static type checker 😁