flax

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