FIX MINOR TUNING BUGS

blackjax-devs · Jun 10, 2024 · ccd9a28 · ccd9a28
1 parent 5ca1b22
commit ccd9a28
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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)
             (