Docs: Fix typo and clarify introduction in Functional API section#5157

copybara-service[bot] merged 2 commits intogoogle:maingoogle/flax:mainfrom

Moriyuki-S:docs/fix-nnx-functional-api-typoMoriyuki-S/flax:docs/fix-nnx-functional-api-typoCopy head branch name to clipboard