Improve integration with external samplers

[juanitorduz]

Motivated by #2124 and specifically by #2124 (comment), we suggest a cleaner API to integrate with other samplers. This includes MCLMC from Blackjax. This is illustrated in the notebook.

(most from the code changes come from the notebook)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[juanitorduz]

Thank you for the feedback @fehiepsi . Let me know how this iteration looks. I am happy to keep iterating until we get the best level of abstraction :)

[reubenharry]

Would people be happy merging this? Certainly it looks good from my end

[juanitorduz]

I would 😄!

[fehiepsi]

Let's wait a bit to rethink whether the design choices are reasonable. Looking like we are changing other internal codes due to the change. It indicates that we will likely break current user code. Let's see if we can come up with simpler solutions.

[juanitorduz]

Great! Always open for feedback 🤗

[juanitorduz]

@fehiepsi you were right, there was no need to add such a large change to the internal API 🙈 . In a463116, I simplified the approach and minimized the changes. Now this PR is way less invasive 🤗

fyi @reubenharry