flax

0521b88c - Merge pull request #2465 from google:proper-scikit-learn-dependency