flax

69233b5f - Merge pull request #1502 from juntang-zhuang/update_optim_doc