def loop_tqdm(n: int, message: typing.Optional[str] = None) -> typing.Callable:
    """
    tqdm progress bar for a JAX fori_loop
    Parameters
    ----------
    n : int
        Number of iterations.
    message : str
        Optional string to prepend to tqdm progress bar.
    Returns
    -------
    typing.Callable:
        Progress bar wrapping function.
    """

    _update_progress_bar, close_tqdm = build_tqdm(n, message)

    def _loop_tqdm(func):
        """
        Decorator that adds a tqdm progress bar to `body_fun`
        used in `jax.lax.fori_loop`.
        """
        def wrapper_progress_bar(i, val):
            _update_progress_bar(i)
            result = func(i, val)
            return close_tqdm(result, i)

        return wrapper_progress_bar

    return _loop_tqdm

    tqdm progress bar for a JAX fori_loop
    Parameters
    ----------
    n : int
        Number of iterations.
    message : str
        Optional string to prepend to tqdm progress bar.
    Returns
    -------
    typing.Callable:
        Progress bar wrapping function.
    
cola/utils/jax_tqdm.py