无法降级至jaxlib==0.1.67且程序运行报错求助

一、当前运行报错详情 #

使用jaxlib==0.1.75时程序抛出ImportError，完整错误栈如下：

from haiku import data_structures
  File "/net/home/justen/.local/lib/python3.10/site-packages/haiku/data_structures.py", line 17, in <module>
    from haiku._src.data_structures import to_immutable_dict
  File "/net/home/justen/.local/lib/python3.10/site-packages/haiku/_src/data_structures.py", line 30, in <module>
    from haiku._src import utils
  File "/net/home/justen/.local/lib/python3.10/site-packages/haiku/_src/utils.py", line 24, in <module>
    import jax
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/__init__.py", line 108, in <module>
    from .experimental.maps import soft_pmap
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/experimental/maps.py", line 25, in <module>
    from .. import numpy as jnp
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/numpy/__init__.py", line 16, in <module>
    from . import fft
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/numpy/fft.py", line 17, in <module>
    from jax._src.numpy.fft import (
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/_src/numpy/fft.py", line 19, in <module>
    from jax import lax
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/lax/__init__.py", line 334, in <module>
    from jax._src.lax.parallel import (
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/_src/lax/parallel.py", line 36, in <module>
    from jax._src.numpy import lax_numpy
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 51, in <module>
    from jax import ops
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/ops/__init__.py", line 16, in <module>
    from jax._src.ops.scatter import (
  File "/net/home/justen/.local/lib/python3.10/site-packages/jax/_src/ops/scatter.py", line 31, in <module>
    from typing import EllipsisType
ImportError: cannot import name 'EllipsisType' from 'typing' (/usr/lib/python3.10/typing.py)

二、环境对比与降级失败情况 #

• 两台机器环境一致：Python 3.9.9虚拟环境，pip==20.0.2，haiku和typing版本完全相同
• 唯一差异：正常运行机器使用jaxlib0.1.67，当前机器使用jaxlib0.1.75
• 尝试降级jaxlib时出现包找不到的错误：

ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.67 (from versions: 0.1.75, 0.1.76, 0.3.0, 0.3.2, 0.3.5, 0.3.7, 0.3.10, 0.3.14, 0.3.15)
ERROR: No matching distribution found for jaxlib==0.1.67

• 尝试指定jax发布源安装仍失败：pip install jaxlib==0.1.67 -f https://storage.googleapis.com/jax-releases/jax_releases.html

三、解决思路 #

• 匹配jax与jaxlib版本：jax和jaxlib版本必须严格对应，先查看正常机器的jax版本，安装对应版本的jax后，再尝试安装jaxlib==0.1.67
• 手动下载whl包本地安装：找到对应Python版本、系统架构的jaxlib==0.1.67的whl文件，使用本地pip命令安装，例如：pip install ./jaxlib-0.1.67-cp39-none-linux_x86_64.whl（需替换为对应自己环境的whl文件名）
• 升级pip版本：旧版pip（20.0.2）可能无法正确识别旧版本的jaxlib包索引，先执行pip install --upgrade pip升级到较新版本，再重新尝试安装jaxlib==0.1.67
• 复制正常机器的jaxlib包：直接从正常机器的site-packages目录复制jaxlib相关文件夹到当前机器的对应虚拟环境目录下，之后可执行pip check验证依赖一致性
• 检查Python版本兼容性：确认jaxlib==0.1.67是否支持Python3.9.9，若存在兼容性问题，可尝试调整虚拟环境的Python版本至与包兼容的版本

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