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]

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().

[cgarciae]

What is this for?

[samanklesaria]

It's a utility datatype used by build_tree_from_paths, which is an inverse of jax.tree.leaves_with_path. It's useful for defining the bridge code porting old rng state to the new names. See the changes to the tutorial.