Progress Bar?

[ahwillia]

Does anyone have any advice on how to display progress as samples are being drawn? I'm currently using the inference loop below (which I copy-pasted from the example notebooks) and would ideally like to have a progress bar that goes from 1 to num_samples during the scan.

def inference_loop(rng_key, kernel, initial_state, num_samples):
    @jax.jit
    def one_step(state, rng_key):
        state, _ = kernel(rng_key, state)
        return state, state

    keys = jax.random.split(rng_key, num_samples)
    _, states = jax.lax.scan(one_step, initial_state, keys)

    return states

[rlouf]

See discussion in #219