flax

9a41e049 - Merge branches 'master' and 'master' of https://github.com/google/flax into rtd-training