From ccd9a284a239f104a2104197565ab868f2c69cfc Mon Sep 17 00:00:00 2001 From: = Date: Tue, 11 Jun 2024 00:01:40 +0200 Subject: [PATCH] FIX MINOR TUNING BUGS --- blackjax/adaptation/adjusted_mclmc_adaptation.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/blackjax/adaptation/adjusted_mclmc_adaptation.py b/blackjax/adaptation/adjusted_mclmc_adaptation.py index 26944a9fe..47df58d2f 100644 --- a/blackjax/adaptation/adjusted_mclmc_adaptation.py +++ b/blackjax/adaptation/adjusted_mclmc_adaptation.py @@ -260,7 +260,9 @@ def L_step_size_adaptation(state, params, num_steps, rng_key): L=params.L * change, step_size=params.step_size * change ) if diagonal_preconditioning: - params = params._replace(sqrt_diag_cov=jnp.sqrt(variances), L = jnp.sqrt(dim)) + params = params._replace( + sqrt_diag_cov=jnp.sqrt(variances), L=jnp.sqrt(dim) + ) initial_da, update_da, final_da = dual_averaging_adaptation(target=target) (