def while_loop_winfo(errorfn, tol, max_iters=None, every=1, desc='', pbar=False, **kwargs):
    """ 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.

        Args:
            errorfn: function of the while state that returns a scalar tracking the error
             (e.g. residual)
            tol: tolerance for errorfn
            every: update progress bar every this many iterations
            desc: description for progress bar
            pbar: whether to show progress bar
            kwargs: additional info to pass to progress bar

        Returns: (tuple) while_loop, info_dict
        """
    info = {}

    def new_while(cond_fun, body_fun, init_val):

        default_desc = f"Running {body_fun.__name__}"

        def construct_info(*_):
            info.pop('errstart', None)
            info.update({'progval': 0, 'errors': [], **kwargs})
            info['iteration_time'] = time.time()
            if pbar:
                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_info(ival):
            i, arg = ival
            error = float(errorfn(arg))
            info['errors'].append(error)
            if pbar:
                errstart = info.setdefault('errstart', error)
                howclose = np.log(error / errstart) / np.log(tol / errstart)
                if max_iters is not None:
                    howclose = max((i + 1) / max_iters, howclose)
                progress = float(min(100 - info['progval'], max(100 * howclose - info['progval'], 0)))
                if progress > 0:
                    info['progval'] += progress
                    info['pbar'].update(progress)

        def close_info(arg):
            i, val = arg
            update_info(arg)
            info['iteration_time'] = (time.time() - info['iteration_time']) / (i + 1)
            if pbar:
                info['pbar'].close()
                info.pop('errstart')
                info.pop('pbar')
            info.pop('progval')
            info['errors'] = np.array(info['errors'][1:])
            info['iterations'] = i + 1

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

        def newcond(ival):
            _, val = ival
            out = jax.lax.cond(
                cond_fun(val),
                lambda _: True,
                lambda _: False,  # host_callback.id_tap(close_info, ival, result=False),
                operand=None)
            return out

        callback(construct_info, None)
        k, val = jax.lax.while_loop(newcond, newbody, (0, init_val))
        callback(close_info, (k, val))
        return val

    return new_while, info
 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.

        Args:
            errorfn: function of the while state that returns a scalar tracking the error
             (e.g. residual)
            tol: tolerance for errorfn
            every: update progress bar every this many iterations
            desc: description for progress bar
            pbar: whether to show progress bar
            kwargs: additional info to pass to progress bar

        Returns: (tuple) while_loop, info_dict
        
cola/utils/jax_tqdm.py