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

PiperOrigin-RevId: 442982471

google/flax,google/flax