flax

3484cda9 - Merge branch 'main' of https://www.github.com/google/flax into howto-colab