wrapper for progress _bar_scan

blackjax/adaptation/window_adaptation.py

@@ -28,7 +28,7 @@
     dual_averaging_adaptation,
 )
 from blackjax.base import AdaptationAlgorithm
-from blackjax.progress_bar import progress_bar_scan
+from blackjax.progress_bar import gen_scan_fn
 from blackjax.types import Array, ArrayLikeTree, PRNGKey
 from blackjax.util import pytree_size
 
@@ -333,23 +333,16 @@ def run(rng_key: PRNGKey, position: ArrayLikeTree, num_steps: int = 1000):
 
         if progress_bar:
             print("Running window adaptation")
-            one_step_ = jax.jit(progress_bar_scan(num_steps)(one_step))
-            start_state = ((init_state, init_adaptation_state), -1)
-        else:
-            one_step_ = jax.jit(one_step)
-            start_state = (init_state, init_adaptation_state)
-
+        scan_fn = gen_scan_fn(num_steps, progress_bar=progress_bar)
+        start_state = (init_state, init_adaptation_state)
         keys = jax.random.split(rng_key, num_steps)
         schedule = build_schedule(num_steps)
-        last_state, info = jax.lax.scan(
-            one_step_,
+        last_state, info = scan_fn(
+            one_step,
             start_state,
             (jnp.arange(num_steps), keys, schedule),
         )
 
-        if progress_bar:
-            last_state, _ = last_state
-
         last_chain_state, last_warmup_state, *_ = last_state
 
         step_size, inverse_mass_matrix = adapt_final(last_warmup_state)

blackjax/progress_bar.py

@@ -94,3 +94,17 @@ def wrapper_progress_bar(carry, x):
         return wrapper_progress_bar
 
     return _progress_bar_scan
+
+
+def gen_scan_fn(num_samples, progress_bar, print_rate=None):
+    if progress_bar:
+
+        def scan_wrap(f, init, *args, **kwargs):
+            func = progress_bar_scan(num_samples, print_rate)(f)
+            carry = (init, -1)
+            (last_state, _), output = lax.scan(func, carry, *args, **kwargs)
+            return last_state, output
+
+        return scan_wrap
+    else:
+        return lax.scan

blackjax/util.py

@@ -11,7 +11,7 @@
 from jax.tree_util import tree_leaves
 
 from blackjax.base import SamplingAlgorithm, VIAlgorithm
-from blackjax.progress_bar import progress_bar_scan
+from blackjax.progress_bar import gen_scan_fn
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 
 
@@ -225,18 +225,12 @@ def one_step(average_and_state, xs, return_state):
     one_step = jax.jit(partial(one_step, return_state=return_state_history))
 
     xs = (jnp.arange(num_steps), keys)
-    if progress_bar:
-        one_step = progress_bar_scan(num_steps)(one_step)
-        (((_, average), final_state), _), history = lax.scan(
-            one_step,
-            (((0, expectation(transform(initial_state))), initial_state), -1),
-            xs,
-        )
-
-    else:
-        ((_, average), final_state), history = lax.scan(
-            one_step, ((0, expectation(transform(initial_state))), initial_state), xs
-        )
+    scan_fn = gen_scan_fn(num_steps, progress_bar)
+    ((_, average), final_state), history = scan_fn(
+        one_step,
+        ((0, expectation(transform(initial_state))), initial_state),
+        xs,
+    )
 
     if not return_state_history:
         return average, transform(final_state)