AttributeError: module 'jax.random' has no attribute 'KeyArray'

[Marsolo1]

Hi,

It seems that the depencencies indicated in the setup.pyare broken.

To reproduce the error:

• Install jax-md in a new virtual environment:

pip install jax-md

Then in a Python interpreter:

from jax_md import rigid_body

Normally you will see the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/martial/vivarium/.venv2/lib/python3.11/site-packages/jax_md/__init__.py", line 27, in <module>
    from jax_md import rigid_body
  File "/Users/martial/vivarium/.venv2/lib/python3.11/site-packages/jax_md/rigid_body.py", line 76, in <module>
    KeyArray = random.KeyArray
               ^^^^^^^^^^^^^^^
  File "/Users/martial/vivarium/.venv2/lib/python3.11/site-packages/jax/_src/deprecations.py", line 53, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.random' has no attribute 'KeyArray'

Thanks in advance for the help

[MarcBerneman]

I will take care of this error. For now, I suggest that you use an older version of JAX.