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)?
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
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)
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'