flax

77b21e2b - Merge pull request #281 from mohitreddy1996:fix-mnist-documentation