flax

5f0ac50c - Merge pull request #2881 from chiamp:rmslayernorm