Flax NNX and Orbax Checkpointing require hacking to work together

I'm building a system using `flax.nnx` and `orbax.checkpointing`. However, it is overly complicated on how to save and restore models due to the new `jax.random.key()` being used in `flax.nnx` rather than `jax.random.PRNGkey()`.

I have had to create a workaround where all layers with `rng` and `key` in their path are changed from `dtype=key` to a format appropriate for saving. Then, upon restoration, they need to be shanged back.

I am attaching a link to a notebook explaining what I've done but I would be keen to hear if there are simpler workarounds? Or, preferably, if there is a way to simple save and restore models?

https://colab.research.google.com/drive/1ozln9ejG7eRtxvbkqHYU3K6OyPvveH9w?usp=sharing

Note: I am also adding an issue to orbax to see if there is a fix their side ([#1337](https://github.com/google/orbax/issues/1337)).