flax

fe995740 - Merge pull request #3783 from google:nnx-jit-accepts-many-Modules