第2章：环境搭建——安装JAX（CPU/GPU/TPU版本）、验证安装、配置开发环境
说实话，很多同学学JAX卡在第一步——环境装不上。我见过太多人花了两天装环境，结果跑demo时报错，心态直接崩了。这一章，我带你一次搞定。
2.1 安装前的准备工作
先看看你的硬件。JAX支持三种后端：CPU、GPU、TPU。我建议你至少有个GPU，不然很多性能调优的乐趣体会不到。
| 后端类型 | 推荐场景 | 硬件要求 |
|---|---|---|
| CPU | 学习入门、小规模实验 | 任何x86_64机器 |
| GPU | 模型训练、大规模计算 | NVIDIA GPU + CUDA 11.8/12.x |
| TPU | Google Cloud专用 | Cloud TPU VM |
我个人习惯用Python 3.10+。为什么？因为JAX对3.12的支持在早期版本有坑，3.10最稳。你可以在终端里跑一下：
python --version
# 确保 >= 3.10
2.2 安装JAX（CPU版本）
CPU版本最简单，说白了就是纯Python包。适合先跑通代码逻辑。
pip install --upgrade pip
pip install jax jaxlib
装完后验证一下：
import jax
print(jax.__version__)
print(jax.devices())
# 应该输出 [cpu(id=0)]
嗯，这里要注意：如果你看到的是空列表，说明安装有问题。我曾经遇到过pip源里jaxlib版本不匹配，导致设备检测不到。解决办法是换官方源：
pip install jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html
2.3 安装JAX（GPU版本）
GPU版本才是重头戏。说白了，JAX的GPU加速是它最大的卖点。你想想看，同样的代码，GPU上跑比CPU快几十倍。
先确认你的CUDA版本：
nvcc --version
# 或者
nvidia-smi # 看右上角的CUDA Version
然后根据CUDA版本选择安装命令：
# CUDA 12.x
pip install --upgrade "jax[cuda12]" jaxlib
# CUDA 11.8
pip install --upgrade "jax[cuda11]" jaxlib
验证GPU是否生效：
import jax
print(jax.devices())
# 应该输出 [gpu(id=0), ...]
⚠️ 避坑指南
我曾经在CUDA 11.8的机器上装了cuda12的jaxlib，结果jax.devices()一直返回空。折腾了两小时才发现是版本不匹配。记住：jaxlib的CUDA版本必须和系统CUDA版本一致。
我曾经在CUDA 11.8的机器上装了cuda12的jaxlib，结果jax.devices()一直返回空。折腾了两小时才发现是版本不匹配。记住：jaxlib的CUDA版本必须和系统CUDA版本一致。
2.4 安装JAX（TPU版本）
TPU版本比较特殊。你需要在Google Cloud上创建TPU VM，然后SSH进去安装。我个人建议初学者先跳过这个，除非你手头有TPU配额。
# 在TPU VM上执行
pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
验证：
import jax
print(jax.devices())
# 应该输出 [tpu(id=0), ...]
2.5 配置VS Code + Jupyter
环境装好了，开发工具也得跟上。我推荐VS Code + Jupyter插件，这套组合拳打起来很顺手。
步骤很简单：
- 安装VS Code（官网下载）
- 安装Python扩展（ms-python.python）
- 安装Jupyter扩展（ms-toolsai.jupyter）
- 在VS Code里新建一个.ipynb文件
然后写个简单的JAX代码测试：
import jax.numpy as jnp
from jax import grad
def f(x):
return x**2 + 3*x + 1
print(grad(f)(2.0)) # 应该输出 7.0
为什么会输出7.0？因为f'(x)=2x+3，在x=2时就是7。JAX的自动微分就是这么丝滑。
💡 我的小技巧
在VS Code里，按Ctrl+Shift+P，输入"Python: Select Interpreter"，选择你装了JAX的那个虚拟环境。不然你可能会发现import jax报错——嗯，我刚开始也犯过这个错。
在VS Code里，按Ctrl+Shift+P，输入"Python: Select Interpreter"，选择你装了JAX的那个虚拟环境。不然你可能会发现import jax报错——嗯，我刚开始也犯过这个错。
2.6 知识体系总览
下面这张图，是我梳理的本章知识结构。你跟着这个流程走，基本不会迷路。
2.7 常见问题与解决方案
我整理了几个高频问题，都是我在项目里踩过的坑：
- 问题：pip install jax报错
原因：网络问题或Python版本太低。解决办法：换国内镜像源pip install -i https://pypi.tuna.tsinghua.edu.cn/simple jax jaxlib - 问题：jax.devices()返回空列表
原因：jaxlib版本与CUDA不匹配。解决办法：卸载重装，指定CUDA版本。 - 问题：VS Code里import jax失败
原因：没选对Python解释器。解决办法：Ctrl+Shift+P，选虚拟环境。
📌 核心要点
环境搭建这件事，说白了就是「版本匹配」四个字。JAX版本、jaxlib版本、CUDA版本、Python版本，任何一个对不上，都可能翻车。我建议你每次装完都跑一遍验证代码，别等到写项目了才发现环境有问题。
环境搭建这件事，说白了就是「版本匹配」四个字。JAX版本、jaxlib版本、CUDA版本、Python版本，任何一个对不上，都可能翻车。我建议你每次装完都跑一遍验证代码，别等到写项目了才发现环境有问题。
好了，环境搭好了，接下来你就可以真正开始玩JAX了。记住：装环境只是开始，后面的性能调优才是重头戏。