flax

551eaf02 - Merge pull request #3989 from google:nnx-fix-vmap