UPDATE TUNING

blackjax-devs · Jun 11, 2024 · 7e21f43 · 7e21f43
1 parent 1ef4c95
commit 7e21f43
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/blackjax/adaptation/adjusted_mclmc_adaptation.py b/blackjax/adaptation/adjusted_mclmc_adaptation.py
@@ -309,6 +309,6 @@ def step(state, key):
         flat_samples = jax.vmap(lambda x: ravel_pytree(x)[0])(samples)
         ess = effective_sample_size(flat_samples[None, ...])
 
-        return state, params._replace(L=(0.4 * params.L) / ess)
+        return state, params._replace(L=(0.4 * params.L) / jnp.mean(ess))
 
     return adaptation_L