4
基础环境
标题
2024-03-28 17:59:38
160

请问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仅这个版本,所以不得不装这个版本。
请问有解决办法或者任何建议吗,非常感谢!

全部评论(1)
只看楼主
Z
zhuwenwen
1楼

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('.'))
2024/3/29
没有更多内容了
--
楼主
Level.1
0 帖子
楼主热帖
暂无数据