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

[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]

thanks @ciguaran!