MapElites notebook dependency errors

Hi,
I'm getting the following error when running the first cell of mapelites.ipynb on Google Colab:
Successfully installed chex-0.1.5 jumanji-0.2.2
---------------------------------------------------------------------------

ModuleNotFoundError                       Traceback (most recent call last)

[<ipython-input-1-7a593bdc590c>](https://localhost:8080/#) in <cell line: 22>()
     22 try:
---> 23     import jumanji
     24 except:

ModuleNotFoundError: No module named 'jumanji'

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)

6 frames

[<ipython-input-1-7a593bdc590c>](https://localhost:8080/#) in <cell line: 22>()
     24 except:
     25     get_ipython().system('pip install "jumanji==0.2.2"')
---> 26     import jumanji
     27 
     28 try:

[/usr/local/lib/python3.10/dist-packages/jumanji/__init__.py](https://localhost:8080/#) in <module>
     13 # limitations under the License.
     14 
---> 15 from jumanji.env import Environment
     16 from jumanji.environments.logic.rubiks_cube import generator as rubik_generator
     17 from jumanji.registration import make, register, registered_environments

[/usr/local/lib/python3.10/dist-packages/jumanji/env.py](https://localhost:8080/#) in <module>
     18 from typing import Any, Generic, Tuple, TypeVar
     19 
---> 20 import chex
     21 from typing_extensions import Protocol
     22 

[/usr/local/lib/python3.10/dist-packages/chex/__init__.py](https://localhost:8080/#) in <module>
     15 """Chex: Testing made fun, in JAX!"""
     16 
---> 17 from chex._src.asserts import assert_axis_dimension
     18 from chex._src.asserts import assert_axis_dimension_comparator
     19 from chex._src.asserts import assert_axis_dimension_gt

[/usr/local/lib/python3.10/dist-packages/chex/_src/asserts.py](https://localhost:8080/#) in <module>
     24 from unittest import mock
     25 
---> 26 from chex._src import asserts_internal as _ai
     27 from chex._src import pytypes
     28 import jax

[/usr/local/lib/python3.10/dist-packages/chex/_src/asserts_internal.py](https://localhost:8080/#) in <module>
     32 
     33 from absl import logging
---> 34 from chex._src import pytypes
     35 import jax
     36 from jax.experimental import checkify

[/usr/local/lib/python3.10/dist-packages/chex/_src/pytypes.py](https://localhost:8080/#) in <module>
     23 ArrayBatched = jax.interpreters.batching.BatchTracer
     24 ArrayNumpy = np.ndarray
---> 25 ArraySharded = jax.interpreters.pxla.ShardedDeviceArray
     26 # Use this type for type annotation. For instance checking,  use
     27 # `isinstance(x, jax.DeviceArray)`.

AttributeError: module 'jax.interpreters.pxla' has no attribute 'ShardedDeviceArray'
adaptive-intelligent-robotics / QDax

MapElites notebook dependency errors #159
