How can I solve my errors and run successfully

[nandinisonar]

Hello Sir/ Ma'am,

I'm writing to request assistance with running JAX-MD in a Colab notebook. I've successfully installed the Jax-MD libraries, but I'm encountering errors.

Could you please help me troubleshoot the errors and guide me on successfully running JAX-MD in a Colab notebook environment?

If possible, could you also advise on the specific error message I'm encountering (once I provide it)?

Thank you for your time and support.

Error in crystal.ipynb

gnome.ipynb

util.ipynb

ModuleNotFoundError Traceback (most recent call last)

in <cell line: 23>()
21 from jax.nn import initializers
22 import jax.numpy as jnp
---> 23 from jax_md import energy
24 from jax_md import partition
25 from jax_md import util

1 frames

/usr/local/lib/python3.10/dist-packages/jax_md/init.py in
13 # limitations under the License.
14
---> 15 from jax_md import space
16 from jax_md import energy
17 from jax_md import minimize

/usr/local/lib/python3.10/dist-packages/jax_md/space.py in
33 from typing import Callable, Union, Tuple, Any
34
---> 35 from jax.abstract_arrays import ShapedArray
36
37 from jax import eval_shape

ModuleNotFoundError: No module named 'jax.abstract_arrays'