Update tutorial: Visualize JAX model

[pavithraes]

This PR updates the Visualize JAX model metrics with TensorBoard tutorial with:

• Seperate sections for better navigation
• Uses the TensorBoard notebook extension for easier reading flow
• General updates to narrative elements

[jakevdp]

Thanks!

[jakevdp]

There are some lint errors – you should be able to fix this by running pre-commit run jupytext --all-files on your branch.

[pavithraes]

Thank you, done!