flax

2b257b45 - Merge pull request #3798 from google:nnx-improve-grad