Test the samplers on Flax

[jeremiecoullon]

Test whether it's easy to train a NN built in Flax using some of the sgmcmc algorithms. Add a notebook in the documentation with the following:

• Build a simple NN (feedforward and/or CNN) using Flax
• train it using a few of the samplers
• Plot the loss curves for these samplers.

[jeremiecoullon]

Added this to documentation