Add NNX support for legacy `jax.random.PRNGKey()`

Currently, it doesn't seem possible to straightforwardly checkpoint (with Orbax) an NNX module that includes random keys (like with dropout), see https://github.com/google/orbax/issues/1105#issue-2488663946. This seems to be due to the new JAX random key type (https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html), which is used here (https://github.com/google/flax/blob/fc19c5dcc136b17d7a2c406831126d93bbdac6d7/flax/nnx/rnglib.py#L186) in NNX. Although Orbax has added individual support for the new type (see https://github.com/google/orbax/issues/620), saving `nnx.state(model)` that includes `dtype=key` doesn't seem to be possible.