Migrate progress bar from fastprogress to tqdm

[zaxtax]

Moves to use tqdm along with adding support for multiple progress bars

Makes blackjax suitable for running multiple chains in parallel.

[codecov]

Codecov Report

Attention: Patch coverage is 91.30435% with 2 lines in your changes are missing coverage. Please review.

Project coverage is 98.80%. Comparing base (7cf4f9d) to head (b2d2273).
Report is 1 commits behind head on main.

❗ Current head b2d2273 differs from pull request most recent head bf703cd. Consider uploading reports for the commit bf703cd to get more accurate results

Files	Patch %	Lines
blackjax/progress_bar.py	91.30%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #655      +/-   ##
==========================================
- Coverage   98.87%   98.80%   -0.07%     
==========================================
  Files          59       59              
  Lines        2745     2752       +7     
==========================================
+ Hits         2714     2719       +5     
- Misses         31       33       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[junpenglao]

from jax.debug import callback?

[junpenglao]

IIUC in the usage we need to specify the num_chains for pmap to work properly. Could you explain a bit more how you are planning to change the API for downstream application so that part works?

blackjax/blackjax/util.py

Line 198 in 7cf4f9d

one_step = progress_bar_scan(num_steps)(_one_step)

blackjax/blackjax/adaptation/window_adaptation.py

Line 329 in 7cf4f9d

one_step_ = jax.jit(progress_bar_scan(num_steps)(one_step))

[zaxtax]

I'm not fully committed to this API, but I was thinking something where along with passing an array of iteration numbers, you also pass in the chain you are currently in. I think this is better than the numpyro design where you are using regexes on device objects to guess what chain to put the computation on.

def inference_loop(rng_key, kernel, initial_state, chain, num_samples, num_chains):

    def _one_step(state, xs):
        _, _, rng_key = xs
        state, _ = kernel(rng_key, state)
        return state, state
    one_step = jax.jit(progress_bar_factory(num_samples, num_chains)(_one_step))

    keys = jax.random.split(rng_key, num_samples)
    _, states = jax.lax.scan(
        one_step,
        initial_state,
        (np.arange(num_samples), chain * np.ones(num_samples), keys),
    )

    return states
    
inference_loop_multiple_chains = jax.pmap(
    inference_loop,
    in_axes=(0, None, 0, 0, None, None),
    static_broadcasted_argnums=(1, 4, 5),
    devices=jax.devices(),
)

[zaxtax]

For downstream applications that don't use multiple chains, I have included logic to maintain backward compatibility. Though I'm not sure how actual code is implementing progress bars for multiple chains today.

[junpenglao]

I see, could you share a small jupyter notebook how it looks like?

[zaxtax]

This is a self-contained example

https://gist.github.com/zaxtax/5fd7c881c6ac83a7ca2798d0a7e230b7

[junpenglao]

Thanks, that's very helpful. Let me think about it a bit.

[zaxtax]

Perfectly happy to rework the API. This is was an attempt to make something simple and backwards compatible.