flax

1e0b9ea5 - Use fork_rngs with implicit state in randomness guide