When importing jumanji, it seems that by default the chex module is pretty old (0.1.3) and the source of this error:
module 'jax.interpreters.pxla' has no attribute 'ShardedDeviceArray'
However, raising chex and optionally jumanji to the latest didn't work. Would you mind trying to look for a solution and reporting it to jumanji/chex if you think it's out of your scope?
I'm trying to run the current https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mapelites.ipynb
When importing jumanji, it seems that by default the
chexmodule is pretty old (0.1.3) and the source of this error:However, raising chex and optionally jumanji to the latest didn't work. Would you mind trying to look for a solution and reporting it to jumanji/chex if you think it's out of your scope?