diff --git a/pymc_experimental/inference/smc/sampling.py b/pymc_experimental/inference/smc/sampling.py index 99488c85..898db598 100644 --- a/pymc_experimental/inference/smc/sampling.py +++ b/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'") diff --git a/tests/test_blackjax_smc.py b/tests/test_blackjax_smc.py index 49db7de7..e0187bd6 100644 --- a/tests/test_blackjax_smc.py +++ b/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