Acme框架运行实验遇JAX属性错误，求依赖兼容解决方案

我正在使用Acme框架开展实验，已按照文档完成安装，但出现源于JAX和HAIKU的AttributeError，目前在GitHub问题中未找到解决方案，恳请帮忙排查依赖问题。

虚拟环境配置 #

dm-acme                      0.4.0
dm-control                   0.0.364896371
dm-env                       1.6
dm-haiku                     0.0.10
dm-launchpad                 0.5.0
dm-reverb                    0.7.0
dm-tree                      0.1.8
acme                         2.10.0
dm-acme                      0.4.0
jax                          0.4.26
jaxlib                       0.4.26+cuda12.cudnn89
python -V                    Python 3.9.5

初始错误详情 #

File "/data/acme/examples/baselines/rl_discrete/run_dqn.py", line 18,
in 
from acme.agents.jax import dqn   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/agents/jax/dqn/__init__.py",
line 18, in 
from acme.agents.jax.dqn.actor import behavior_policy   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/agents/jax/dqn/actor.py",
line 20, in 
from acme.agents.jax import actor_core as actor_core_lib   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/actor_core.py",
line 22, in 
from acme.jax import networks as networks_lib   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/networks/__init__.py",
line 18, in 
from acme.jax.networks.atari import AtariTorso   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/networks/atari.py",
line 29, in 
from acme.jax.networks import base   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/networks/base.py",
line 24, in 
import haiku as hk   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/haiku/__init__.py",
line 20, in 
from haiku import experimental   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/haiku/experimental/__init__.py", line 34, in 
from haiku._src.dot import abstract_to_dot   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/haiku/_src/dot.py",
line 163, in 
@jax.linear_util.transformation   File "/data/acme/acme_venv_new/lib/python3.9/site-packages/jax/_src/deprecations.py",
line 54, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}") AttributeError: module 'jax' has no attribute 'linear_util'

该错误似乎源于haiku与JAX，请问如何修复？

更新尝试 #

根据建议重新安装了jax和jaxlib，但又出现新错误：

Traceback (most recent call last):
  File "/data/acme/examples/baselines/rl_discrete/run_dqn.py", line 18, in <module>
    from acme.agents.jax import dqn
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/agents/jax/dqn/__init__.py", line 18, in <module>
    from acme.agents.jax.dqn.actor import behavior_policy
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/agents/jax/dqn/actor.py", line 20, in <module>
    from acme.agents.jax import actor_core as actor_core_lib
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/actor_core.py", line 22, in <module>
    from acme.jax import networks as networks_lib
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/networks/__init__.py", line 45, in <module>
    from acme.jax.networks.multiplexers import CriticMultiplexer
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/networks/multiplexers.py", line 20, in <module>
    from acme.jax import utils
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/acme/jax/utils.py", line 190, in <module>
    devices: Optional[Sequence[jax.xla.Device]] = None,
  File "/data/acme/acme_venv_new/lib/python3.9/site-packages/jax/_src/deprecations.py", line 53, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax' has no attribute 'xla'

完整pip freeze列表已整理成文档，同时查阅了相关GitHub问题，针对此jax.xla错误有何更新建议或可行的解决办法？谢谢

内容的提问来源于stack exchange，提问作者Jerry07