SMC-MCMC integration test, plus fixes.

[ciguaran]

Thank you for opening a PR!

Some kernel interfaces weren't compatible with SMC. This PR fixes them and adds a test that protects us from regressions that may break the integration between the MCMC kernels and SMC. If we believe a kernel could be used within SMC, then it should be added to this test.

[codecov]

Codecov Report

Merging #522 (0583dad) into main (a650f9b) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #522   +/-   ##
=======================================
  Coverage   99.28%   99.28%           
=======================================
  Files          47       47           
  Lines        1948     1948           
=======================================
  Hits         1934     1934           
  Misses         14       14           

Impacted Files	Coverage Δ
blackjax/kernels.py	99.21% <100.00%> (ø)
blackjax/mcmc/random_walk.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[albcab]

The changes to the random walk algorithms look fine to me. Only test_kernel_compatibility.py seems frivolous, we are simply checking that the low level kernel function has a certain order in its inputs: kernel(key, state, logprob_fn, *kernel_params). Feels like more tests simply to check this are unnecessary, @junpenglao @rlouf thoughts?

[ciguaran]

The changes to the random walk algorithms look fine to me. Only test_kernel_compatibility.py seems frivolous, we are simply checking that the low level kernel function has a certain order in its inputs: kernel(key, state, logprob_fn, *kernel_params). Feels like more tests simply to check this are unnecessary, @junpenglao @rlouf thoughts?

Although the fix was indeed changing the order of the inputs, this test is checking that a step SMC step can be run with certain MCMC inner kernel. That implies correct order of inputs, but also that the output of the inner kernel has certain structure (for example, returns state, info, the state has a .position method, etc.).

The point is that before this PR we were able to break the integration between MCMC and SMC, and no test would complain. By adding this test we are including a 'safety net' that alerts the developer that a new change is breaking. The fact that the test is simple is actually a feature, not a bug, since it will be very clear from the test output which regression was included.

[albcab]

The point is that before this PR we were able to break the integration between MCMC and SMC, and no test would complain. By adding this test we are including a 'safety net' that alerts the developer that a new change is breaking.

But the developer of a new algorithm would still need to manually add a test that checks for this, its not done automatically. For example NUTS, MALA, MGrad, GHMC, elliptical slice are not checked by this safety net, could still break it.

I see your point though, maybe we can add support (in the test) for NUTS and MALA and use this as a safety net for the popular algorithms.

[ciguaran]

The point is that before this PR we were able to break the integration between MCMC and SMC, and no test would complain. By adding this test we are including a 'safety net' that alerts the developer that a new change is breaking.

But the developer of a new algorithm would still need to manually add a test that checks for this, its not done automatically. For example NUTS, MALA, MGrad, GHMC, elliptical slice are not checked by this safety net, could still break it.

Yes, new algorithms need to be added as they are coded. Only if we expect the user to be able to use them within SMC, it would be part of correctly testing the new algorithm.

I see your point though, maybe we can add support (in the test) for NUTS and MALA and use this as a safety net for the popular algorithms.

Added NUTS and MALA