flax

Docs: Fix typo and clarify introduction in Functional API section

#5157

Merged