flax

88011afe - Add PRNG design notes to help new JAX users, refactor the train_epoch explanation