Fix Blackjax SMC integration (#374)

Author: ciguaran

pymc_experimental/inference/smc/sampling.py

@@ -25,6 +25,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+from blackjax.smc import extend_params
 from blackjax.smc.resampling import systematic
 from pymc import draw, modelcontext, to_inference_data
 from pymc.backends import NDArray
@@ -126,16 +127,20 @@ def sample_smc_blackjax(
 
     if kernel == "HMC":
         mcmc_kernel = blackjax.mcmc.hmc
-        mcmc_parameters = dict(
-            step_size=inner_kernel_params["step_size"],
-            inverse_mass_matrix=jnp.eye(posterior_dimensions),
-            num_integration_steps=inner_kernel_params["integration_steps"],
+        mcmc_parameters = extend_params(
+            dict(
+                step_size=inner_kernel_params["step_size"],
+                inverse_mass_matrix=jnp.eye(posterior_dimensions),
+                num_integration_steps=inner_kernel_params["integration_steps"],
+            )
         )
     elif kernel == "NUTS":
         mcmc_kernel = blackjax.mcmc.nuts
-        mcmc_parameters = dict(
-            step_size=inner_kernel_params["step_size"],
-            inverse_mass_matrix=jnp.eye(posterior_dimensions),
+        mcmc_parameters = extend_params(
+            dict(
+                step_size=inner_kernel_params["step_size"],
+                inverse_mass_matrix=jnp.eye(posterior_dimensions),
+            )
         )
     else:
         raise ValueError(f"Invalid kernel {kernel}, valid options are 'HMC' and 'NUTS'")

tests/test_blackjax_smc.py

@@ -80,7 +80,6 @@ def fast_model():
         ("NUTS", False, {"step_size": 0.1}),
     ],
 )
-@pytest.mark.xfail(reason="Still need to investigate")
 def test_sample_smc_blackjax(kernel, check_for_integration_steps, inner_kernel_params):
     """
     When running the two gaussians model