flax

27c06b62 - Avoid assert_array_equal on PRNG keys.

Avoid assert_array_equal on PRNG keys. This operates via conversion to np.array, which will soon be disallowed by https://github.com/jax-ml/jax/pull/24481.