[Python] Is there a way to monitor progress of vmapped function in JAX (like with tqdm)?

Let's say I have an expensive jitted function f(x) and I've vmapped it over a large array of inputs -- is there any way to monitor the progress of how many of the vmapped inputs are done being processed? For example in normal python without JAX we would use something like tqdm (or print statements):

MWE:

import jax
from tqdm import tqdm

def f(x):
return x**2

xarr = jnp.arange(1000)

# with JAX, I just call a vectorized f on the input xarr
# but if f was expensive, I want to know how many x in xarr have been processed
out = jax.jit(jax.vmap(f))(xarr)

# if we were in non-JAX python, we would normally use tqdm to print a progress bar
out2 = []
for i in tqdm(range(len(xarr))):
out2.append(f(xarr))

# or instead of tqdm, i could just print but we can't with jitted/traced jax functions
print('finished %i/%i'%(i,len(xarr)))

Continue reading...