Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Aug 15, 2024
1 parent 3b65bb1 commit 0aff4f6
Showing 1 changed file with 12 additions and 7 deletions.
19 changes: 12 additions & 7 deletions pymc_experimental/inference/smc/sampling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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'")
Expand Down

0 comments on commit 0aff4f6

Please sign in to comment.