Blackjax SMC from pymc models

[ciguaran]

Allows to sample Pymc's models using Blackjax's Sequential Monte Carlo implementations. Apart from getting a Jax-based implementation for Blackjax, this PR allows for using HMC and NUTS as kernels, which aren't available in the existing PyMC implementation of SMC. Moreover, diagnosis are exposed and stored in the resulting arviz.InferenceData

In order to sample using BJ SMC, we need

• code to compute jaxified logprior and loglikelihood functions, applyable over SMC Particles.
• code to compute diagnosis over the SMC run.
• code to build an arviz.InferenceData object from the sampler's output.

[ciguaran]

@junpenglao it seems that blackjax is not included in windows-environment-test.yml although is present in environment-test.yml do you know the reason for that?

[ricardoV94]

@ciguaran because JAX isn't/wasn't compatible with Windows

[junpenglao]

Maybe add also a random walk metropolis so we can compare with the SMC in PyMC easier?
Also NUTS-SMC is not very well tested IIUC.

[ciguaran]

That's the idea, but it has been difficult to use the algorithm without any tuning on the inner kernel parameters between SMC runs. I am planning on merging the optimization PR on blackjax first and then use it from here. The SMC in Pymc does apply those optimizations.

[ciguaran]

@ciguaran because JAX isn't/wasn't compatible with Windows

gotcha, it seems they have an experimental version for windows, I've added it in the reqs to see what happens, otherwise we can remove the test from the windows suite if you agree @ricardoV94

[jessegrabowski]

Can you add a notebook example showing how to use this?

[aloctavodia]

Just a few changes to the docstrings and this is good to go.

[aloctavodia]

thanks @ciguaran!