remove unused input in scan

blackjax-devs · Oct 23, 2023 · ae4ec10 · ae4ec10
1 parent 51cf08a
commit ae4ec10
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/blackjax/adaptation/window_adaptation.py b/blackjax/adaptation/window_adaptation.py
@@ -297,7 +297,7 @@ def window_adaptation(
     )
 
     def one_step(carry, xs):
-        _, rng_key, adaptation_stage = xs
+        rng_key, adaptation_stage = xs
         state, adaptation_state = carry
 
         new_state, info = mcmc_kernel(
@@ -335,7 +335,7 @@ def run(rng_key: PRNGKey, position: ArrayLikeTree, num_steps: int = 1000):
         last_state, info = jax.lax.scan(
             one_step_,
             (init_state, init_adaptation_state),
-            (jnp.arange(num_steps), keys, schedule),
+            (keys, schedule),
         )
         last_chain_state, last_warmup_state, *_ = last_state