UPDATE STREAMING AVG

blackjax-devs · May 19, 2024 · 9dd740f · 9dd740f
1 parent 0ab0694
commit 9dd740f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/blackjax/adaptation/mclmc_adaptation.py b/blackjax/adaptation/mclmc_adaptation.py
@@ -460,10 +460,10 @@ def step(iteration_state, weight_and_key):
             adaptive_state = adaptive_state._replace(log_step_size=jnp.log(step_size))
             # step_size = 1e-3
 
+            x = ravel_pytree(state.position)[0]
             # update the running average of x, x^2
             streaming_avg = streaming_average(
-                O=lambda x: jnp.array([x, jnp.square(x)]),
-                x=ravel_pytree(state.position)[0],
+                expectation=jnp.array([x, jnp.square(x)]),
                 streaming_avg=streaming_avg,
                 weight=(1 - mask) * success * step_size,
                 zero_prevention=mask,