flax

0f95fae6 - Merge branch 'main' of https://www.github.com/google/flax into scope-tests