flax

376804f3 - Merge pull request #2486 from yotarok:fix_batch_norm_init