flax

2ac765a5 - Removes default_key_shape() from scope.py since jax.random.default_prng_impl is part of the latest JAX release now.