Following the suggestions in the #65 issue, here are some details for how we could add the progress bar.
Following the JAX discussion (which I wrote up in a blog post), we can define a decorator that adds a progress bar to the body_fun used in lax.scan:
def _print_consumer(arg, transform):
iter_num, n_iter = arg
print(f"Iteration {iter_num}/{n_iter}")
@jit
def _progress_bar(arg, result):
"""
Print progress of a scan/loop only if the iteration number is a multiple of the print_rate
Usage: carry = progress_bar((iter_num, n_iter, print_rate), carry)
"""
iter_num, n_iter, print_rate = arg
result = lax.cond(
iter_num % print_rate==0,
lambda _: host_callback.id_tap(_print_consumer, (iter_num, n_iter), result=result),
lambda _: result,
operand = None)
return result
def progress_bar(num_samples):
"""
Decorator that adds a progress bar to `body_fun` used in `lax.scan`.
Note that `body_fun` must be looping over `jnp.arange(num_samples)`.
This means that `iter_num` is the current iteration number
"""
def pbar_factory(func):
print_rate = int(num_samples/10)
def wrapper_progress_bar(carry, iter_num):
iter_num = _progress_bar((iter_num, num_samples, print_rate), iter_num)
return func(carry, iter_num)
return wrapper_progress_bar
return pbar_factory
To use it we simply add the progress_bar decorator to the update_scan function in mcx/sample.py. So the code would become:
The problem is that this progress bar can only work if update_scan has access to the current iteration number. However it loops over the keys rather than jnp.arange(num_samples). I see 2 possible solutions:
We change the code so that lax.scan loops over jnp.arange(num_samples) insteads of the keys, but that means changing more code elsewhere to split the keys at each iteration
We simply scan over both jnp.arange(rng_keys.shape[0]) as well as rng_keys (ie: the JAX version of Python's enumerate!). This might be a much simpler fix and hopefully doesn't change too much.
I'll try out this second option. Any thoughts?
Note that once this works there's still the issue that the progress bar is very simple and isn't as nice as tqdm.
Following the suggestions in the #65 issue, here are some details for how we could add the progress bar.
Following the JAX discussion (which I wrote up in a blog post), we can define a decorator that adds a progress bar to the
body_fun
used inlax.scan
:To use it we simply add the
progress_bar
decorator to theupdate_scan
function in mcx/sample.py. So the code would become:Problem:
The problem is that this progress bar can only work if
update_scan
has access to the current iteration number. However it loops over the keys rather thanjnp.arange(num_samples)
. I see 2 possible solutions:lax.scan
loops overjnp.arange(num_samples)
insteads of the keys, but that means changing more code elsewhere to split the keys at each iterationjnp.arange(rng_keys.shape[0])
as well asrng_keys
(ie: the JAX version of Python'senumerate
!). This might be a much simpler fix and hopefully doesn't change too much.I'll try out this second option. Any thoughts?
Note that once this works there's still the issue that the progress bar is very simple and isn't as nice as tqdm.