运行Haiku时出现jaxlib.xla_extension无PmapFunction属性报错如何解决

报错原因 #

该报错由Haiku和JAX/jaxlib版本不兼容触发：0.4.14及更高版本的jaxlib移除了jaxlib.xla_extension.PmapFunction定义，而未适配该改动的低版本Haiku，在/usr/local/lib/python3.7/dist-packages/haiku/_src/transform.py的check_not_jax_transformed(f)校验逻辑中仍然调用了该废弃属性，最终抛出属性不存在异常。

修复方案 #

• 方案一：升级Haiku到适配版本
  直接执行命令升级到最新稳定版Haiku即可，新版已经适配了高版本jaxlib的改动，替换了对应校验逻辑：
  pip install --upgrade dm-haiku
• 方案二：降级jaxlib到兼容版本
  如果项目依赖特定旧版Haiku无法升级，可将jaxlib降级到0.4.13及之前的版本，注意要保证jax和jaxlib版本对齐，避免触发其他兼容性问题：
  pip install jaxlib==0.4.13 jax==0.4.13
• 方案三：临时修改源码（仅用于测试场景，不推荐生产环境使用）
  打开报错对应的haiku/_src/transform.py文件，找到check_not_jax_transformed(f)函数，删除其中引用jaxlib.xla_extension.PmapFunction的判断分支即可临时绕过校验。

内容的提问来源于stack exchange，提问作者Laura K.