flax

5a2e9c4b - Merge branch 'master' of https://github.com/google/flax into rtd-training