flax

ebfe1152 - [nnx] improve mnist tutorial