flax

cb3f3fa4 - Merge pull request #289 from mohitreddy1996:refactor-mnist