Multiple Progress Bars

[zombie-einstein]

Issue #24 (thanks also to @andrewlesak for the prototype) and some of my own projects have highlighted that it would be useful to be able to print multiple progress bars when the scan or loop is inside a vmap. Currently this will push all the updates to a single progress bar (as the scanned function is compiled once for the vmap, and so only one progress bar is initialised), showing false progress updates when using vmap.

My approach here has been to allow the user to inject an id for the progress bar, which is then carried through the loop or scan, which then allows updates to be routed to individual progress bars. For example with a scan

from jax_tqdm import BarId, scan_tqdm

@scan_tqdm(n, print_rate=10)
def step(carry, _):
    return carry + 1, carry + 1

@jax.jit
def inner(i):
    init = BarId(i=i, carry=0)
    _final, _all_numbers = jax.lax.scan(step, init, jnp.arange(n))
    return _final.carry, _all_numbers

last_numbers, all_numbers = jax.vmap(inner)(jax.numpy.arange(5))

Here BarId is a simple class that contain the id of the progress bar and the carried value. The user needs to wrap the init of the scan in a BarId and provide an integer id, and conversely unpack the results afterwards. Using a class like this is an easy way of detecting when a user has provided a bar-id.

This approach means the original scan function can still be used without modification (you can just provide the carry without wrapping), and overall it has a fairly low impact on the overall API. I guess it adds some complication (mainly arising from needing to initialise the bars at runtime rather than compile time) requiring users to bundle up the scan arguments.

Welcome any thoughts on this API change!

*Note that some terminals don't support updating multiple bars (for example the PyCharm terminal) this is a known bug with TQDM

[jeremiecoullon]

LGTM! 🚀

[zombie-einstein]

LGTM! 🚀

Cheers, I just need to add instructions to the README then I can release.