add u32[4] as possible PRNGKeyArray type

[jondeaton]

jax.random.key(0, impl="rbg") gives

Array((), dtype=key<rbg>) overlaying:
[0 0 0 0]

[patrick-kidger]

Should this not already be handled by the Key dtype? (Which is distinct from PRNGKeyArray, which largely exists for backwards compatibility with now-very-old JAX.)

[jondeaton]

good point Key already seems to cover this. thank you

[jondeaton]

actually this is to cover the case of

jax.random.PRNGKey(0, impl="rbg")
Array([0, 0, 0, 0], dtype=uint32)

[patrick-kidger]

Merged!