tutorial 2 (intro to jax): checkpoint API has changed

[murphyk]

** Tutorial **

https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial2/Introduction_to_JAX.html

Describe the bug

The checkpoint saving/loading code needs to be changed to the new flax API

https://flax.readthedocs.io/en/latest/guides/orbax_upgrade_guide.html

[phlippe]

Thanks for pointing out the API change! It seems to me that the API is still in progress and possible subject to change, so I'll wait a bit longer before changing it everywhere in the tutorials.