flax

61580b99 - Merge pull request #965 from levskaya:optimizer_fixes