Change PRNGKey to key in jax-101 PRNG guide

[IvyZX]

jax.random.PRNGKey use cases should be replaced by jax.random.key as planned by #17297.

[jakevdp]

Hi - I think @froystig has a branch that does this, but we've been waiting on fuller support for new-style keys in downstream libraries before swapping over the docs. @froystig, what's the status of that?

[froystig]

I could rebase and update the docs branch, and then look for places to include a disclaimer like the one you added to the RNG tutorial here. Then we can aim to merge it prior to waiting for more downstream libraries. How about that?

[IvyZX]

What kind of downstream libraries are we waiting for? I can definitely add a disclaimer there, but if the blocker is downstream should we explain that (like, this new API doesn't work for bla bla cases)?

[froystig]

The challenge is that jax itself is backwards compatible by way of polymorphism: all jax.random functions accept both old- and new-style keys, and those that return keys (split and fold_in) do so based on what flavor of key they were given. Meanwhile, other libraries that deal with jax's keys might have not yet upgraded to new-style. If our docs encourage new-style keys to users, then those users might pass new-style keys to a library that doesn't understand the new style yet.

I don't think anything is blocking this really, and my guess is that we only need disclaimers in our own docs for now. Does that sound right to you?

[froystig]

#17741 (which Jake was referring to above) included similar changes to these files and has been merged.