flax

21f500ca - reverting changes on mnist_tutorial.md