Update NNX installation instructions in mnist tutorial

[Davidnet]

What does this PR do?

• Specifies that nnx is now part of flax, and just suggest to install flax with pip

Fixes # (issue)

Checklist

• [x] This PR fixes a minor issue

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Davidnet]

Closing since upgrading flax breaks GPU install version of jax on Colab