flax

2ad97319 - Merge pull request #4183 from google:nnx-fsdp-toy-example