flax

21695fdf - Merge pull request #5157 from Moriyuki-S:docs/fix-nnx-functional-api-typo