Use prng_key from RandomnessState instead of passing it explicitly

[copybara-service[bot]]

Use prng_key from RandomnessState instead of passing it explicitly

This will allow for using NeuralGCM models will less boilerplate.