Skip to content

Commit

Permalink
Change isokinetic_integrator generation API (#689)
Browse files Browse the repository at this point in the history
  • Loading branch information
junpenglao authored Jun 3, 2024
1 parent a3e235f commit 83bc3a0
Showing 1 changed file with 1 addition and 2 deletions.
3 changes: 1 addition & 2 deletions blackjax/mcmc/integrators.py
Original file line number Diff line number Diff line change
Expand Up @@ -374,9 +374,8 @@ def format_isokinetic_state_output(

def generate_isokinetic_integrator(coefficients):
def isokinetic_integrator(
logdensity_fn: Callable, *args, **kwargs
logdensity_fn: Callable, sqrt_diag_cov: ArrayTree = 1.0
) -> GeneralIntegrator:
sqrt_diag_cov = kwargs.get("sqrt_diag_cov", 1.0)
position_update_fn = euclidean_position_update_fn(logdensity_fn)
one_step = generalized_two_stage_integrator(
esh_dynamics_momentum_update_one_step(sqrt_diag_cov),
Expand Down

0 comments on commit 83bc3a0

Please sign in to comment.