def build_tqdm(n: int, message: typing.Optional[str] = None) -> typing.Tuple[typing.Callable, typing.Callable]:
    """
    Build the tqdm progress bar on the host
    """

    if message is None:
        message = f"Running for {n:,} iterations"
    tqdm_bars = {}

    if n > 20:
        print_rate = int(n / 20)
    else:
        print_rate = 1
    remainder = n % print_rate

    def _define_tqdm(arg):
        tqdm_bars[0] = tqdm(range(n))
        tqdm_bars[0].set_description(message, refresh=False)

    def _update_tqdm(arg):
        tqdm_bars[0].update(float(arg))

    def _update_progress_bar(iter_num):
        "Updates tqdm from a JAX scan or loop"
        _ = jax.jax.lax.cond(
            iter_num == 0,
            lambda _: callback(_define_tqdm, None),
            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),
            lambda _: None,
            operand=None,
        )

        _ = jax.lax.cond(
            # update tqdm by `remainder`
            iter_num == n - remainder,
            lambda _: callback(_update_tqdm, remainder),
            lambda _: None,
            operand=None,
        )

    def _close_tqdm(arg):
        tqdm_bars[0].close()
        return arg

    def close_tqdm(result, iter_num):
        return jax.lax.cond(
            iter_num == n - 1,
            lambda _: io_callback(_close_tqdm, result),
            lambda _: result,
            operand=None,
        )

    return _update_progress_bar, close_tqdm

    Build the tqdm progress bar on the host
    
cola/utils/jax_tqdm.py