Update mnist example to use NNX (clone of #5002)#5064

samanklesaria wants to merge 1 commit intogoogle:maingoogle/flax:mainfrom

samanklesaria:pr/5002samanklesaria/flax:pr/5002Copy head branch name to clipboard