Bounded while loop

jax_tqdm/__init__.py

@@ -1 +1 @@
-from jax_tqdm.pbar import PBar, loop_tqdm, scan_tqdm
+from jax_tqdm.pbar import PBar, bounded_while_tqdm, loop_tqdm, scan_tqdm

jax_tqdm/pbar.py

@@ -10,8 +10,9 @@
 
 @chex.dataclass
 class PBar:
-    id: int
     carry: typing.Any
+    id: int = 0
+    iter: int = 0
 
 
 def scan_tqdm(
@@ -51,7 +52,7 @@ def _scan_tqdm(func):
         """
 
         def wrapper_progress_bar(carry, x):
-            if type(x) is tuple:
+            if isinstance(x, tuple):
                 iter_num, *_ = x
             else:
                 iter_num = x
@@ -62,11 +63,11 @@ def wrapper_progress_bar(carry, x):
                 carry, x = update_progress_bar((carry, x), iter_num, bar_id=bar_id)
                 result = func(carry, x)
                 result = (PBar(id=bar_id, carry=result[0]), result[1])
-                return close_tqdm(result, iter_num, bar_id=bar_id)
+                return close_tqdm(result, n, iter_num, bar_id=bar_id)
             else:
                 carry, x = update_progress_bar((carry, x), iter_num)
                 result = func(carry, x)
-                return close_tqdm(result, iter_num)
+                return close_tqdm(result, n, iter_num)
 
         return wrapper_progress_bar
 
@@ -115,17 +116,49 @@ def wrapper_progress_bar(i, val):
                 i, val = update_progress_bar((i, val), i, bar_id=bar_id)
                 result = func(i, val)
                 result = PBar(id=bar_id, carry=result)
-                return close_tqdm(result, i, bar_id=bar_id)
+                return close_tqdm(result, n, i, bar_id=bar_id)
             else:
                 i, val = update_progress_bar((i, val), i)
                 result = func(i, val)
-                return close_tqdm(result, i)
+                return close_tqdm(result, n, i)
 
         return wrapper_progress_bar
 
     return _loop_tqdm
 
 
+def bounded_while_tqdm(
+    cond_fun: typing.Callable,
+    body_fun: typing.Callable,
+    n: int,
+    print_rate: typing.Optional[int] = None,
+    tqdm_type: str = "auto",
+    **kwargs,
+) -> typing.Tuple[typing.Callable, typing.Callable]:
+
+    update_progress_bar, close_tqdm = build_tqdm(n, print_rate, tqdm_type, **kwargs)
+
+    def cond_fun_wrapper(val: PBar) -> bool:
+        return cond_fun(val.carry)
+
+    def close_bar(val, iter_num, bar_id):
+        return close_tqdm(val, iter_num, iter_num - 1, bar_id=bar_id)
+
+    def cont(val, _iter_num, _bar_id):
+        return val
+
+    def body_fun_wrapper(val: PBar) -> PBar:
+        iter_num = val.iter
+        bar_id = val.id
+        val = val.carry
+        val = update_progress_bar(val, iter_num, bar_id=bar_id)
+        val = body_fun(val)
+        val = jax.lax.cond(cond_fun(val), close_bar, cont, val, iter_num, bar_id)
+        return PBar(carry=val, id=bar_id, iter=iter_num + 1)
+
+    return cond_fun_wrapper, body_fun_wrapper
+
+
 def build_tqdm(
     n: int,
     print_rate: typing.Optional[int],
@@ -178,9 +211,6 @@ def build_tqdm(
                 f"number of steps {n}, got {print_rate}"
             )
 
-    remainder = n % print_rate
-    remainder = remainder if remainder > 0 else print_rate
-
     def _define_tqdm(bar_id: int):
         bar_id = int(bar_id)
         tqdm_bars[bar_id] = pbar(
@@ -193,9 +223,9 @@ def _define_tqdm(bar_id: int):
     def _update_tqdm(bar_id: int):
         tqdm_bars[int(bar_id)].update(print_rate)
 
-    def _close_tqdm(bar_id: int):
+    def _close_tqdm(bar_id: int, final_value: int):
         _pbar = tqdm_bars.pop(int(bar_id))
-        _pbar.update(remainder)
+        _pbar.update(int(final_value) - _pbar.n)
         _pbar.clear()
         _pbar.close()
 
@@ -214,22 +244,24 @@ def _inner_update(i, _carry):
             )
             return _carry
 
+        cond = iter_num > 0
+
         carry = jax.lax.cond(
-            iter_num == 0,
-            _inner_init,
+            cond,
             _inner_update,
+            _inner_init,
             iter_num,
             carry,
         )
 
         return carry
 
-    def close_tqdm(result: typing.Any, iter_num: int, bar_id: int = 0):
+    def close_tqdm(result: typing.Any, target: int, iter_num: int, bar_id: int = 0):
         def _inner_close(_result):
-            callback(_close_tqdm, bar_id, ordered=True)
+            callback(_close_tqdm, bar_id, target, ordered=True)
             return _result
 
-        result = jax.lax.cond(iter_num + 1 == n, _inner_close, lambda r: r, result)
+        result = jax.lax.cond(iter_num + 1 == target, _inner_close, lambda r: r, result)
         return result
 
     return update_progress_bar, close_tqdm

tests/test_examples.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 import pytest
 
-from jax_tqdm import PBar, loop_tqdm, scan_tqdm
+from jax_tqdm import PBar, bounded_while_tqdm, loop_tqdm, scan_tqdm
 
 
 @pytest.mark.parametrize("print_rate", [None, 1, 10])
@@ -61,6 +61,23 @@ def inner(i):
     assert jnp.array_equal(all_numbers, jnp.tile(1 + jnp.arange(n), (n_maps, 1)))
 
 
+def test_bounded_while_loop():
+    n_total = 10_000
+    n_stop = 5_000
+
+    def cond_fun(x):
+        return x < n_stop
+
+    def body_fun(x):
+        return x + 1
+
+    cond_fun, body_fun = bounded_while_tqdm(cond_fun, body_fun, n=n_total)
+    init_val = PBar(carry=0)
+    result = jax.lax.while_loop(cond_fun, body_fun, init_val)
+
+    assert result == 5_000
+
+
 @pytest.mark.parametrize("print_rate", [None, 1, 10])
 def test_vmap_w_loop(print_rate):
     n = 10_000