From 0aff4f626819f02dbfd1b3c2686adbaa600d7384 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Aug 2024 20:16:40 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pymc_experimental/inference/smc/sampling.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/pymc_experimental/inference/smc/sampling.py b/pymc_experimental/inference/smc/sampling.py index 9c011841c..898db598b 100644 --- a/pymc_experimental/inference/smc/sampling.py +++ b/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'")