pip安装jax/jaxlib后导入报错：CPU不支持AVX指令

问题原因 #

你通过pip安装的预编译jaxlib是针对支持AVX指令集的现代CPU构建的，而你的旧款x86_64 CPU不支持AVX指令集，因此触发兼容性报错。

解决方案 #

1. 安装无AVX依赖的预编译jaxlib版本 #

JAX官方提供了适配无AVX CPU的预编译包，执行以下命令安装对应版本：

pip install jaxlib==0.4.23 jax==0.4.23 --find-links https://storage.googleapis.com/jax-releases/jaxlib_releases.html

该命令会自动匹配你的Python版本和系统环境，下载并安装不依赖AVX的jaxlib包。

2. 临时跳过CPU特性检查（不推荐） #

如果仅需临时测试功能，可以通过设置环境变量或JAX配置跳过检查，但此方式可能导致运行时崩溃或性能大幅下降：

• 方式一：运行脚本前设置环境变量

  JAX_DISABLE_JIT=1 python /mod/CreateModels.py
  

• 方式二：在代码开头添加配置

  import os
  os.environ['JAX_DISABLE_JIT'] = '1'
  import jax
  

3. 从源码编译jaxlib（适合需要自定义优化的场景） #

如果预编译包无法满足需求，可以从源码编译兼容旧CPU的jaxlib：

• 安装编译依赖（以Debian/Ubuntu为例）：

  sudo apt update && sudo apt install build-essential cmake python3-dev
  

• 从JAX官方源码仓库克隆代码，进入仓库目录后执行编译：

  cd jax
  python build/build.py --cpu_features=generic
  

• 安装编译好的wheel包：

  pip install dist/jaxlib-*.whl
  

• 确保jax版本与jaxlib匹配：

  pip install jax==0.4.23
  

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