Python3.6 conda环境导入haiku遇AttributeError：jax.tree_util无PyTreeDef属性

在Python3.6的conda环境中结合Jax使用Haiku时，出现以下导入错误，更新Jax版本后问题仍未解决：

Traceback (most recent call last):
  File "train.py", line 14, in <module>
    import haiku as hk
  File "/home/young/.local/lib/python3.6/site-packages/haiku/__init__.py", line 19, in <module>
    from haiku import data_structures
  File "/home/young/.local/lib/python3.6/site-packages/haiku/data_structures.py", line 18, in <module>
    from haiku._src.data_structures import to_haiku_dict
  File "/home/young/.local/lib/python3.6/site-packages/haiku/_src/data_structures.py", line 176, in <module>
    class FlatComponents(NamedTuple):
  File "/home/young/.local/lib/python3.6/site-packages/haiku/_src/data_structures.py", line 178, in FlatComponents
    structure: jax.tree_util.PyTreeDef
AttributeError: module 'jax.tree_util' has no attribute 'PyTreeDef'

修复方案 #

• 升级Python版本到3.7及以上：Python3.6已停止官方支持，Jax从0.2.26版本开始不再兼容Python3.6，而jax.tree_util.PyTreeDef是Jax较新版本才引入的属性，Haiku的新版本依赖该属性，导致Python3.6下无法同时满足Jax和Haiku的版本要求。升级到Python3.7+后，直接安装最新版Jax和Haiku即可解决：

  pip install --upgrade jax jaxlib dm-haiku
  

• 若无法升级Python，安装兼容旧版本：指定安装支持Python3.6的Jax和Haiku旧版本，例如：

  pip install jax==0.2.25 jaxlib==0.1.75 dm-haiku==0.0.10
  

  注意jaxlib需匹配Python3.6的系统架构，若找不到对应wheel，可能需要从源码编译或查找第三方提供的兼容包。

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