flax

updated rng docstring for init, apply and make_rng

#3765

Merged