Implement SVGD

[antotocar34]

As discussed in #385.

Done:

• Implementation of the algorithm
• Added very simple (bad) test.

TODO:

• Add an example with GPU where gradients is calculated with pmap.
• How to take mini batch gradients?
• Come up with a better test.

Appreciate any comments & feedback!

A few important guidelines and requirements before we can merge your PR:

• [ X ] If I add a new sampler, there is an issue discussing it already;
• [ X ] We should be able to understand what the PR does from its title only;
• [ X ] There is a high-level description of the changes;
• [ X ] There are links to all the relevant issues, discussions and PRs;
• [ X ] The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• [ X ] pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• [ ~ ] There are tests covering the changes;
• The doc is up-to-date;
• If I add a new sampler* I added/updated related examples

[antotocar34]

I've made a quick collab notebook to play around with the algorithm.

https://colab.research.google.com/drive/115rjkthFjz10d_l_M2BqAQdn-kaDdulC?usp=sharing

[albcab]

Thanks for this PR!

Let me know if there is anything that needs clarification on my review (or if I've missed anything).

To be able to do pmap and to be able to use minbatches for the gradient computations, you can consider taking as input an optional grad_estimator or grad_fn that allows the user to specify its own gradient computations, and have it default to the usual jax.grad of the log density. In the case of minibatches, the script sgmcmc/gradients.py has all the basic components you'll need already implemented!

An example notebook meant to run on GPU can be left as a separate PR for blackjax/sampling-book, it would make a very nice addition to this implementation.

I might be the best use of your time to leave the tests for when you are on the final draft of the implementation, i.e. when you have minibatch gradients, pmapping, etc.

[antotocar34]

Thanks for the detailed comments! I'll try and get to them next week 😃

[codecov]

Codecov Report

Merging #512 (8f29ad3) into main (87d94fe) will decrease coverage by 0.02%.
The diff coverage is 97.72%.

❗ Current head 8f29ad3 differs from pull request most recent head b9e8347. Consider uploading reports for the commit b9e8347 to get more accurate results

@@            Coverage Diff             @@
##             main     #512      +/-   ##
==========================================
- Coverage   99.30%   99.29%   -0.02%     
==========================================
  Files          48       48              
  Lines        2019     1989      -30     
==========================================
- Hits         2005     1975      -30     
  Misses         14       14              

Impacted Files	Coverage Δ
blackjax/__init__.py	100.00% <ø> (ø)
blackjax/kernels.py	99.25% <91.66%> (ø)
blackjax/vi/__init__.py	100.00% <100.00%> (ø)
blackjax/vi/svgd.py	100.00% <100.00%> (ø)

... and 23 files with indirect coverage changes

[albcab]

I had to implement SVGD for a project I'm working on. The latest commit is a working version with the improvements from my review. I've tested it on medium-sized models and the vmapping hasn't broken yet, so we can leave it like this for now. It might still need some work on the documentation, other than that it should be ready to merge.

[junpenglao]

Thank you both!