flax

c84510d0 - Merge pull request #3936 from google:nnx-stabilize