SGLD visualization notebook

[tcz]

I took the liberty of adding a new notebook to the examples. I played around with SGLD visualizations and I thought others might find this useful to help the concept click.

Feel free to leave this unmerged.

You can see the notebook here:
https://colab.research.google.com/drive/1J9titx3OEDSY_Fjq0zfCXB2NlkWFZsWV?usp=sharing

What this notebook contains:

1. A simplified LLC estimation without temperature scaling and localization ($n\beta = 1$, $\gamma = 0$)
2. A linear toy model with a convex bowl loss landscape, SGLD path visualized
3. Four different toy datasets
4. A non-linear model with four parameters with loss landscapes and SGLD paths visualized across all possible combinations of two parameters
5. LLC vs. loss evolution over several checkpoints, with loss landscape and SGLD path visualizations for the first two parameters