Improve and migrate model_surgery guide

[IvyZX]

This PR does the following:

• Add a bunch more texts explaining why and how to do param and opt state tweaking to readers with little context of Flax terminologies.
• Explained the FrozenDict a bit more.
• Make the guide notebook-based, so that it's trivial to open a colab and test the code.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[gnecula]

It is a good idea to transition docs from .rst to Colab+md files. Thanks!

[8bitmp3]

It is a good idea to transition docs from .rst to Colab+md files. Thanks!

• 1 @gnecula We have some docs that use tables with before/after code. Unless we find a Markdown-way of implementing this without any (or too many) (external) Markdown "themes" for Sphinx, we can keep them for now. There may also be a way of doing this in HTML embedded in Markdown 🤔 (we do this in TensorFlow . org)

@IvyZX Thanks for the PR. Will review asap.

[8bitmp3]

It is a good idea to transition docs from .rst to Colab+md files. Thanks!

(Note: In some Markdown flavors, code diff highlighting can be done by adding "diff" (similar to adding "python" if you want the code to have python syntax colors). Will investigate if this works in Sphinx for ReadTheDocs. cc @cgarciae You can use it in GitHub comments that support Markdown (during reviews). Reference: https://stackoverflow.com/questions/40883421/diff-syntax-highlighting-in-github-markdown )

[marcvanzee]

In my opinion, the Github diff highlighting is quite a lot worse than what we have now. We currently have a side-by-side view, which I think is very effective to demonstrate the effect of a change (see e.g. @cgarciae's BatchNorm guide). The diff view with green and red is not so easy to read, I think.

I think these side-by-side views do have a benefit so I would be in favor of keeping them as .rst for now. I do agree that if we don't use these side-by-side views in a doc we should switch to notebooks.

[8bitmp3]

side-by-side views do have a benefit so I would be in favor of keeping them as .rst for now.

+1 If there's no Markdown-flavored way of doing this, we can keep those RST docs that use that view.

[IvyZX]

I agree that the side-by-side views look pretty convenient and we should keep them as .rst, until we find a good alternative. The doc on this PR doesn't have the side-by-side views so it could be safely converted.

[8bitmp3]

After @cgarciae and @levskaya are finished with their reviews, I will add my final suggestions.

[chiamp]

I think the FrozenDict explanation is very helpful, thanks for adding it!