fix

Author: ciguaran

pymc_experimental/inference/smc/sampling.py

@@ -24,11 +24,7 @@
 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
 from pymc.backends import NDArray
@@ -133,14 +129,14 @@ def sample_smc_blackjax(
         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"],
-        ))
+            num_integration_steps=inner_kernel_params["integration_steps"])
+        )
     elif kernel == "NUTS":
         mcmc_kernel = blackjax.mcmc.nuts
         mcmc_parameters = extend_params(dict(
             step_size=inner_kernel_params["step_size"],
-            inverse_mass_matrix=jnp.eye(posterior_dimensions),
-        ))
+            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,10 +17,6 @@
 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
@@ -84,11 +80,8 @@ 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