Fix

Author: ciguaran

pymc_experimental/inference/smc/sampling.py

@@ -24,6 +24,10 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+<<<<<<< Updated upstream
+=======
+from blackjax.smc import extend_params
+>>>>>>> Stashed changes
 
 from blackjax.smc.resampling import systematic
 from pymc import draw, modelcontext, to_inference_data
@@ -126,17 +130,17 @@ def sample_smc_blackjax(
 
     if kernel == "HMC":
         mcmc_kernel = blackjax.mcmc.hmc
-        mcmc_parameters = dict(
+        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(
+        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

@@ -17,6 +17,10 @@
 import pytensor.tensor as pt
 import pytest
 import scipy
+<<<<<<< Updated upstream:tests/test_blackjax_smc.py
+=======
+from blackjax.smc import extend_params
+>>>>>>> Stashed changes:pymc_experimental/tests/test_blackjax_smc.py
 
 from numpy import dtype
 from xarray.core.utils import Frozen
@@ -80,7 +84,11 @@ def fast_model():
         ("NUTS", False, {"step_size": 0.1}),
     ],
 )
+<<<<<<< Updated upstream:tests/test_blackjax_smc.py
 @pytest.mark.xfail(reason="Still need to investigate")
+=======
+
+>>>>>>> Stashed changes:pymc_experimental/tests/test_blackjax_smc.py
 def test_sample_smc_blackjax(kernel, check_for_integration_steps, inner_kernel_params):
     """
     When running the two gaussians model