Repository: pyro-ppl/numpyro

Issue:
Progress bars only update with num_chains=1
Running [Bayesian Regression](http://num.pyro.ai/en/stable/tutorials/bayesian_regression.html) tutorial with the default `num_chains=1`: it updates itself as NumPyro samples:

![1 chain](https://user-images.githubusercontent.com/33491632/123524469-79c6c000-d6c2-11eb-803d-e35a975cd2a8.gif)

With `num_chains=4`: all four progress bars stay at 0/2000:

![4 chains](https://user-images.githubusercontent.com/33491632/123524473-81866480-d6c2-11eb-8d9f-25fc403dd006.gif)



Discussion:
Hi @MarcoGorelli, I think you'll need to install [ipywidgets](https://ipywidgets.readthedocs.io/en/latest/user_install.html) for jupyter notebook/lab. I think we should detect if a user already has `ipywidgets` installed and raise a warning. Do you want to take a stab at this? I think all we need to do is to add an additional logic after [this check](https://github.com/pyro-ppl/numpyro/blob/77be4d8171b0570bcabbebf1456f70ef94f9d0e8/numpyro/util.py#L310)
```python
        if num_chains > 1 and progbar and issubclass(tqdm_auto, tqdm.notebook.tqdm):
            try:
                import ipywidgets
            except ImportError:
                warnings.warn(a warning to tell users to install ipywidgets to display progress bar in gpu)
                progbar = False
            
```
yup, will do!
I just tried the same example on master, installed ipywidgets, and got

```python-traceback
ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 197, in _close_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_3'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 197, in _close_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_0'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 197, in _close_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_2'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 192, in _update_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'

ERROR:absl:Postponing exception raised in callback function: Traceback (most recent call last):
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1583, in _outfeed_receiver_callback
    callback(arrays, device)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 1004, in _outside_call_run_callback
    raise e  # Let the exception propagate
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 947, in _outside_call_run_callback
    res = callback(arg, device, unpacked_transforms)
  File "/home/marco/numpyro-dev/venv/lib/python3.8/site-packages/jax/experimental/host_callback.py", line 606, in <lambda>
    callback = lambda arg, device, transforms: callback_func(arg, transforms, device=device)
  File "/home/marco/numpyro-dev/numpyro/util.py", line 197, in _close_tqdm
    chain = int(str(device)[4:])
ValueError: invalid literal for int() with base 10: '_CPU_1'
```

-----

Sorry to bother you with this, I'll try to diagnose the issue, just writing it here in meantime in case the issue is obvious to you
I think the issue is we have a new pattern for `device` in your system. To be safe, I think we can do
```
chain = int(str(device).split("_")[-1])
```
to capture the device number. What do you think?

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -204,6 +204,11 @@ class MCMC(object):
     .. note:: Setting `progress_bar=False` will improve the speed for many cases. But it might
         require more memory than the other option.
 
+    .. note:: If setting `num_chains` greater than `1` in a Jupyter Notebook, then you will need to
+        have installed `ipywidgets <https://ipywidgets.readthedocs.io/en/latest/user_install.html>`_
+        in the environment from which you launced Jupyter in order for the progress bars to render
+        correctly.
+
     :param MCMCKernel sampler: an instance of :class:`~numpyro.infer.mcmc.MCMCKernel` that
         determines the sampler for running MCMC. Currently, only :class:`~numpyro.infer.hmc.HMC`
         and :class:`~numpyro.infer.hmc.NUTS` are available.
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -21,6 +21,7 @@
 from jax.tree_util import tree_flatten, tree_map
 
 _DISABLE_CONTROL_FLOW_PRIM = False
+_CHAIN_RE = re.compile(r"(?<=_)\d+$")  # e.g. get '3' from 'TFRT_CPU_3'
 
 
 def set_rng_seed(rng_seed):
@@ -189,12 +190,16 @@ def progress_bar_factory(num_samples, num_chains):
         tqdm_bars[chain].set_description("Compiling.. ", refresh=True)
 
     def _update_tqdm(arg, transform, device):
-        chain = int(str(device)[4:])
+        chain_match = _CHAIN_RE.search(str(device))
+        assert chain_match
+        chain = int(chain_match.group())
         tqdm_bars[chain].set_description(f"Running chain {chain}", refresh=False)
         tqdm_bars[chain].update(arg)
 
     def _close_tqdm(arg, transform, device):
-        chain = int(str(device)[4:])
+        chain_match = _CHAIN_RE.search(str(device))
+        assert chain_match
+        chain = int(chain_match.group())
         tqdm_bars[chain].update(arg)
         finished_chains.append(chain)
         if len(finished_chains) == num_chains:


Test patch: