Jax will likely support a stateful PRNG object similar to Flax's RngStreams in the future: jax-ml/jax#28845.
Eventually, we'll want to have flax use these native flax objects. To ease that eventual transition, we should give RngStreams the same interface (specifically, the key and split methods) and express the new explicit-rng docs in terms of this interface.