Docs: Fix typo and clarify introduction in Functional API section

[Moriyuki-S]

What does this PR do?

This PR improves the documentation for the Flax NNX Functional API (docs_nnx/nnx_basics.ipynb).

Specifically, it:

1. Fixes a typo: Corrected "example of of" to "define a".
2. Clarifies the introduction: The original text stated that the example code "uses the Functional API". However, the immediate code block only defines the class and does not yet use split, merge, or update. I updated the text to clarify that we are defining a module to be used in the following sections to demonstrate the API.

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
• This change is discussed in a Github issue/discussion (please add a link).
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests. (No quality testing = no merge!)

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB