Support new JAX custom PRNG key system.
JAX recently introduced an extension of its PRNG Key system to
support alternative software PRNG systems as well as hardware PRNG
systems such as those used in TPUs. This fixes flax's runtime typecheck
to support these new-style keys as well.
PiperOrigin-RevId: 400934167
