replaced seed with key

[Tomas542]

Replaced all seed values with key for uniformity of style with jax.random

[Tomas542]

@rdyro 3 month later, but I've finished it)

[Tomas542]

Also should close issue #1137 after PR got accepted

[carlosgmartin]

@Tomas542 JAX now recommends using random.key over random.PRNGKey, which will eventually be deprecated. I suggest editing this PR to use the former.

[Tomas542]

@carlosgmartin yes, I know, but there are some problems with this in tests where NumPy checks for type compatibility. The new keys are of type key<fry> or something like that. It is possible iteratively just replace some of the old styles with the new ones without rewriting the tests if necessary. Tell me, I'll do it

P.S. Accidentally, I've made a mistake in hungarian_algorithm_test, where there were no such problem.

[rdyro]

The changes look great so far!

@carlosgmartin is right that we should move to random.key and you're right @Tomas542 that it breaks the current hyperparameter validation funtion in optax.

Let me see if we can fix the optax hyperparameter validation before this PR so that we can stick to random.key