bump blackjax in test environment

Author: ciguaran

conda-envs/environment-test.yml

@@ -11,5 +11,5 @@ dependencies:
 - statsmodels
 - pip:
   - pymc>=5.16.1  # CI was failing to resolve
-  - blackjax
+  - blackjax>=1.2.3
   - scikit-learn

pymc_experimental/inference/smc/sampling.py

@@ -24,6 +24,7 @@
 import jax
 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
@@ -126,16 +127,20 @@ def sample_smc_blackjax(
 
     if kernel == "HMC":
         mcmc_kernel = blackjax.mcmc.hmc
-        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"])
+        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 = extend_params(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'")