flax

4c472878 - [nnx] improve Optimizer metadata propagation