make jax callback compatible with tqdm (#20)

* make jax callback compatible with tqdm * ordered callbacks to flush progress bar creation --------- Co-authored-by: Matthew Mould <[email protected]>
jeremiecoullon · May 15, 2024 · cd2fe02 · cd2fe02
1 parent 43d5fda
commit cd2fe02
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/jax_tqdm/pbar.py b/jax_tqdm/pbar.py
@@ -131,29 +131,29 @@ def _define_tqdm(arg, transform):
         tqdm_bars[0].set_description(message, refresh=False)
 
     def _update_tqdm(arg, transform):
-        tqdm_bars[0].update(arg)
+        tqdm_bars[0].update(int(arg))
 
     def _update_progress_bar(iter_num):
         "Updates tqdm from a JAX scan or loop"
         _ = jax.lax.cond(
             iter_num == 0,
-            lambda _: callback(_define_tqdm, None, None),
+            lambda _: callback(_define_tqdm, None, None, ordered=True),
             lambda _: None,
             operand=None,
         )
 
         _ = jax.lax.cond(
             # update tqdm every multiple of `print_rate` except at the end
             (iter_num % print_rate == 0) & (iter_num != n - remainder),
-            lambda _: callback(_update_tqdm, print_rate, None),
+            lambda _: callback(_update_tqdm, print_rate, None, ordered=True),
             lambda _: None,
             operand=None,
         )
 
         _ = jax.lax.cond(
             # update tqdm by `remainder`
             iter_num == n - remainder,
-            lambda _: callback(_update_tqdm, remainder, None),
+            lambda _: callback(_update_tqdm, remainder, None, ordered=True),
             lambda _: None,
             operand=None,
         )
@@ -164,7 +164,7 @@ def _close_tqdm(arg, transform):
     def close_tqdm(result, iter_num):
         _ = jax.lax.cond(
             iter_num == n - 1,
-            lambda _: callback(_close_tqdm, None, None),
+            lambda _: callback(_close_tqdm, None, None, ordered=True),
             lambda _: None,
             operand=None,
         )