def pbar_while(errorfn, tol, desc='', every=1, hide=False):
    """ Decorator for while loop with progress bar. Assumes that
        errorfn is a function of the loop variable and returns a scalar
        that starts at a given value and decreases to tol as the loop progresses."""
    if hide:
        return jax.lax.while_loop

    def new_while(cond_fun, body_fun, init_val):
        info = {'progval': 0, 'pbar': None}
        default_desc = f"Running {body_fun.__name__}"

        def construct_tqdm(arg):
            _bar_format = "{l_bar}{bar}| {n:.3g}/{total_fmt} [{elapsed}<{remaining}, "
            _bar_format += "{rate_fmt}{postfix}]"
            info['pbar'] = tqdm(total=100, desc=f'{desc or default_desc}', bar_format=_bar_format)

        def update_tqdm(arg):
            error = errorfn(arg)
            errstart = info.setdefault('errstart', error)
            progress = max(100 * np.log(error / errstart) / np.log(tol / errstart) - info['progval'], 0)
            progress = min(100 - info['progval'], progress)
            if progress > 0:
                info['progval'] += progress
                info['pbar'].update(float(progress))

        def close_tqdm(arg):
            update_tqdm(arg)
            info['pbar'].close()
            return False

        def newbody(ival):
            i, val = ival
            jax.lax.cond(
                i % every == 0,
                lambda _: callback(update_tqdm, val),
                lambda _: None,
                operand=None,
            )
            return (i + 1, body_fun(val))

        def newcond(ival):
            i, val = ival
            out = jax.lax.cond(cond_fun(val), lambda _: True, lambda _: io_callback(close_tqdm, val), operand=None)
            return out

        callback(construct_tqdm)
        _, val = jax.lax.while_loop(newcond, newbody, (0, init_val))
        return val

    return new_while
 Decorator for while loop with progress bar. Assumes that
        errorfn is a function of the loop variable and returns a scalar
        that starts at a given value and decreases to tol as the loop progresses.
cola/utils/jax_tqdm.py