flax

2cfd1747 - Merge pull request #3947 from google:nnx-vmap-improve-axes-size-detection