zombie-einstein

Hey @sash-a and @Chulabhaya, Yeah the defaults are pretty good, when I was testing with [Mava](https://github.com/instadeepai/Mava) I was using these settings (with the rest of the defaults used): ``` name:...

There's some notes in the [original PR for the environment](https://github.com/instadeepai/jumanji/pull/259#issuecomment-2604387430), may help. Though apologies they are not super detailed!

Thanks @Chulabhaya glad you are using it. I should have also said [I forked Mava](https://github.com/zombie-einstein/Mava) when I was testing, this may also help you out. I am planning to implement...

Do you have a simple code example? i don't think there is an immediate work around (especially since we somewhat rely on having ordered debug events) but it would be...

Hi @jwtkeeble, apologies I just noticed I didn't actually reply to this (several months ago)! I did try your example, and don't think there is a simple fix for this,...

Hi @jwtkeeble, I think there may actually be an issue with `pmap`, this example ``` import os os.environ["XLA_FLAGS"] = ( '--xla_force_host_platform_device_count=4 ' # simulate with 4 CPUs ) import jax...

Looking into this a bit more, there is a flag we set to ensure ordered events in a few places (e.g. [here](https://github.com/jeremiecoullon/jax-tqdm/blob/b1b611dbccbdf60c1cfda70ba1d92df05374d7a7/jax_tqdm/pbar.py#L206)). Setting this to False and a modifications to...

Definitely +1 for vmap, I recently ran into a use-case for this and was going to look into a solution, would be interested to see how you got this working.

Are you suggesting we add a wrapper around an existing bounded while loop implementation?

I see what you are doing with the vmap, would be interested in adding similar functionality. May need to think about how to implement it as part of the API...