Sampling for Numpyro breaks running in parallel show progress mode

When I try to run MCMC for parallel chains in new version the sampling for chains sequentially works so does parallel sampling without progress. However with progress enabled it throws error. Any help will be appreciated.

Maybe this is due to not having ipywidgets, from the MCMC docs:

If setting num_chains greater than 1 in a Jupyter Notebook, then you will need to have installed ipywidgets in the environment from which you launced Jupyter in order for the progress bars to render correctly.

Could you check if Relax matching device pattern by fehiepsi · Pull Request #1101 · pyro-ppl/numpyro · GitHub resolves the issue? We will make a patch release if it fixes the issue.

I have the ipywidgets installed. As proof, you can see the progress bar is rendered correctly for multiple chains in sequential runs and actually progress bar shows up in parallel run too.

It works. thanks for looking into it.