“Jax”的版本间差异

来自Shiyin's note
跳到导航 跳到搜索
无编辑摘要
无编辑摘要
 
(未显示同一用户的1个中间版本)
第4行: 第4行:
==安装==
==安装==
*jax和jaxlib版本要匹配(注意cuda支持)
*jax和jaxlib版本要匹配(注意cuda支持)
*我的CUDA版本是11.1 ,卡是A40,系统是Ubuntu,cudnn版本是805
*我的CUDA版本是11.1 ,卡是A40,系统是[[Ubuntu]],cudnn版本是805
*安装最新版本 jax 0.2.26和jaxlib0.1.75后会在random函数报错“CustomCall failed: jaxlib/cuda_prng_kernels.cc:30: operation cudaGetLastError() failed: the provided PTX was compiled with an unsupported toolchain”
*安装最新版本 jax 0.2.26和jaxlib0.1.75后会在random函数报错“CustomCall failed: jaxlib/cuda_prng_kernels.cc:30: operation cudaGetLastError() failed: the provided PTX was compiled with an unsupported toolchain”
*最后安装的 jax0.2.1 (pip install -v jax==0.2.1), jaxlib是0.1.71 [https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.71+cuda111-cp38-none-manylinux2010_x86_64.whl]解决问题
:*最后安装的 jax0.2.2 (pip install -v jax==0.2.2), jaxlib是0.1.72 [https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.72+cuda111-cp38-none-manylinux2010_x86_64.whl]解决问题
:*好像是因为一定要11.1的驱动11的不行(或者cudnn的问题)

2023年5月11日 (四) 13:28的最新版本

  • 简单的说就是GPU加速、支持自动微分(autodiff)的numpy。参考[1]
  • [2]

安装

  • jax和jaxlib版本要匹配(注意cuda支持)
  • 我的CUDA版本是11.1 ,卡是A40,系统是Ubuntu,cudnn版本是805
  • 安装最新版本 jax 0.2.26和jaxlib0.1.75后会在random函数报错“CustomCall failed: jaxlib/cuda_prng_kernels.cc:30: operation cudaGetLastError() failed: the provided PTX was compiled with an unsupported toolchain”
  • 最后安装的 jax0.2.2 (pip install -v jax==0.2.2), jaxlib是0.1.72 [3]解决问题
  • 好像是因为一定要11.1的驱动11的不行(或者cudnn的问题)