Add a documentation page on checkpointing

[IvyZX]

This page is meant to address the following:

• Explain how to save and load a checkpoint
• What to do if you don't know the original object, what to do if it can be one of multiple objects (e.g., a TrainState with or without batch_stats)
• How to use it in multiprocess

Live preview: https://flax--2530.org.readthedocs.build/en/2530/guides/use_checkpointing.html

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov-commenter]

Codecov Report

Merging #2530 (041f5e0) into main (a479e13) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #2530   +/-   ##
=======================================
  Coverage   79.41%   79.41%           
=======================================
  Files          49       49           
  Lines        5188     5188           
=======================================
  Hits         4120     4120           
  Misses       1068     1068           

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[cgarciae]

Hey @IvyZX, to enable jupytext for this notebook run:

jupytext --set-formats ipynb,md:myst docs/guides/use_checkpointing.ipynb

More info here:

https://flax.readthedocs.io/en/latest/advanced_topics/contributing.html#creating-new-notebooks

[8bitmp3]

Update in IvyZX#3