flax

c6498d57 - [nnx] add Optimizer keep_model param