flax

7bcee15f - Merge pull request #4470 from marcelroed:nnx-optimizer-typing