Unsolvable jax dependency

[adiaconu11]

Following the most recent commit, the jax.tree module is now being used instead of jax.tree_*. This is change requires jax >= 0.4.25. However, the are still many parts of the repository that are still old/deprecated. For instance, if you install 0.4.25 you might get something like:

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

This is because this module has been removed in jax 0.4.24, meaning that in order to not run into this problem you need jax <=0.4.23. Obviously this goes against the requirement above.

Lastly, there is still the issue with DeviceArray and ShardedDeviceArray. They have all been changed to somply jax.Array back in jax=0.4.0! At the current state of the repo you basically need to add lines like:

jax.interpreters.xla.DeviceArray = jax.Array in order to be able to even import acme...

[adiaconu11]

the jax.Array issue can be solved by installing chex==0.1.7 instead of the default 0.1.6. This should be updated in the requirements of acme.

[lkoelman]

Bumping into similar problems. It is very hard to install a functional stack.

[fload]

same couldn't get it installed in windows or Ubuntu

[TaroKingCN]

same couldn't get it installed in windows or Ubuntu

have you solved it ?

[AryanVBW]

This is a comprehensive summary of the JAX compatibility issues. The core problem is that Acme's codebase uses deprecated JAX APIs that have been removed in newer versions.

Working installation approach (as of late 2024):

# Create clean environment
python -m venv acme-env
source acme-env/bin/activate

# Install compatible JAX version first
pip install jax==0.4.20 jaxlib==0.4.20

# Install chex 0.1.7 (fixes jax.Array issues)
pip install chex==0.1.7

# Install acme with legacy resolver
pip install git+https://github.com/deepmind/acme.git#egg=dm-acme[jax,tf,envs] --use-deprecated=legacy-resolver

Key version constraints:

• jax <= 0.4.23 (avoids KeyArray removal in 0.4.24)
• jax >= 0.4.0 (for basic compatibility)
• chex == 0.1.7 (fixes DeviceArray/ShardedDeviceArray → jax.Array migration)

The underlying issues:

1. jax.random.KeyArray removed in JAX 0.4.24
2. jax.xla.Device deprecated, replaced with jax.Device
3. DeviceArray/ShardedDeviceArray replaced with jax.Array
4. jax.tree_* functions moved to jax.tree module

Note: These are codebase updates that would need to be made to Acme itself for proper support of JAX >= 0.4.25. Until then, pinning to older versions is the only reliable workaround.