flax

b626099b - Merge pull request #3759 from google:nnx-improve-graph-update