Matthew Mould

I think the following is related. On jax 0.4.28 I get `KeyError: 0` in `_update_tqdm` from `tqdm_bars[0].update(arg)`, even when iterations begin at 0.

> > I think the following is related. On jax 0.4.28 I get `KeyError: 0` in `_update_tqdm` from `tqdm_bars[0].update(arg)`, even when iterations begin at 0. > > I've actually been...

> I wonder if the `ordered` flag [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.callback.html#jax-debug-callback)is what we are looking for? Good spot! I just added `ordered=True` to all calls of `callback` and my previous failure cases seem...

This is a bit late, but another option is to defined individual bijections for the non-transformed variables and the remaining ones, and then stack them together into a single bijection:...