Add split and key methods to RngStream and Rngs

[samanklesaria]

Closes #5046

[samanklesaria]

Unfortunately, the key attribute on streams already exists. I have renamed the existing attribute key_, and replaced its use with the hand sed command sed -I '' "s/\.key\([^_(]\)/.key_\1/".

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[cgarciae]

Thanks @samanklesaria ! I'll have to run all internal tests before merging.

As a follow up PR, we could consider:

• replace all usage of jax.random.<function>(rngs(), ...) with rngs.<function>(...). I've done some of this.
• when we do need a key, favor rngs.key() or rngs.some_stream.key().