flax

4a15be78 - [flax:benchmarks] Fix flax benchmarks.

[flax:benchmarks] Fix flax benchmarks. - Getting build OOMs (inline some things to avoid tf dependency). - Update mnist to nnx. - Make each benchmark its own target. PiperOrigin-RevId: 872413358