请问https://download.sourcefind.cn:65024/4/main/jax/dtk22.10里的jaxlib-0.1.72有其对应的jax包吗?
我尝试过的解决方法:
直接使用jax-0.1.72会报错:
Traceback (most recent call last):
File "/path/to/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py", line 156, in <module>
import jax
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/__init__.py", line 17, in <module>
from .api import (
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/api.py", line 38, in <module>
from . import core
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/core.py", line 31, in <module>
from . import dtypes
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/dtypes.py", line 31, in <module>
from .lib import xla_client
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/lib/__init__.py", line 34, in <module>
version = tuple(int(x) for x in jaxlib_version.__version__.split('.'))
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/lib/__init__.py", line 34, in <genexpr>
version = tuple(int(x) for x in jaxlib_version.__version__.split('.'))
ValueError: invalid literal for int() with base 10: '72+git6c419e3'
就算把这里的这个version直接赋值为(0,1,72),也会报错:
Traceback (most recent call last):
File "/path/to/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py", line 156, in <module>
import jax
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/__init__.py", line 17, in <module>
from .api import (
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/api.py", line 38, in <module>
from . import core
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/core.py", line 31, in <module>
from . import dtypes
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/dtypes.py", line 31, in <module>
from .lib import xla_client
File "/path/to/pdeenv/lib/python3.9/site-packages/jax/lib/__init__.py", line 54, in <module>
from jaxlib import pytree
ImportError: cannot import name 'pytree' from 'jaxlib' (/path/to/pdeenv/lib/python3.9/site-packages/jaxlib/__init__.py)
项目依赖的pde求解包要求python3.9,且资源中适配python3.9的jaxlib仅这个版本,所以不得不装这个版本。
请问有解决办法或者任何建议吗,非常感谢!

暂无数据
pip install jax==0.2.21,/path/to/pdeenv/lib/python3.9/site-packages/jax/lib/__init__.py修改为以下两行即可:
version = jaxlib_version.__version__.split('+')
version = tuple(int(x) for x in version[0].split('.'))