第2章环境搭建

第2章：环境搭建——安装JAX（CPU/GPU/TPU版本）、验证安装、配置开发环境

说实话，环境搭建这事儿，看着简单，但坑是真不少。我见过太多人卡在这一步，代码写好了，结果装不上，心态直接崩了。咱们今天就把这事儿捋清楚，一次搞定。

2.1 安装前的准备工作

先看看你的机器。JAX 对 Python 版本有要求，我个人习惯用 Python 3.10 或 3.11，太老或太新都可能出问题。你可以在终端里敲一下：

python --version

如果版本不对，建议用 conda 或 pyenv 管理环境。我一般这么干：

conda create -n jax_env python=3.11
conda activate jax_env

嗯，这里要注意：别在 base 环境里直接装，隔离环境是基本素养。

2.2 安装JAX——CPU版本

CPU 版本最简单，适合先跑通流程。直接 pip 就行：

pip install --upgrade pip
pip install jax jaxlib

装完之后，验证一下：

import jax
print(jax.__version__)
print(jax.devices())

如果输出类似 [TFRT_CPU_0]，说明成了。我曾经在 MacBook 上装完，发现它默认用了 CPU，心里还挺踏实——至少能跑。

2.3 安装JAX——GPU版本

GPU 版本稍微麻烦点。你得先确认 CUDA 和 cuDNN 版本。JAX 对 CUDA 版本有严格限制，比如 JAX 0.4.30 要求 CUDA 12.2 以上。

我的建议是：先查一下官方兼容性表。然后执行：

pip install --upgrade pip
pip install jax[cuda12] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

注意，这里用的是 cuda12，如果你的是 CUDA 11，就换成 cuda11。装完后验证：

import jax
print(jax.devices())

如果看到 [CUDA_0] 之类的，恭喜你，GPU 跑起来了。我曾经在服务器上装完，发现它死活认不出 GPU，最后发现是 cuDNN 版本不对——嗯，这种坑踩过一次就记住了。

⚠️ 避坑指南： 我曾经因为 CUDA 版本和 JAX 不匹配，折腾了整整一下午。后来学乖了，每次装之前先看官方文档的兼容性表格。别偷懒，这一步省不了。

2.4 安装JAX——TPU版本

TPU 版本主要用在 Google Cloud 上。如果你有 TPU 虚拟机，可以这样装：

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

验证方式一样：

import jax
print(jax.devices())

输出应该是 [TPU_0, TPU_1, ...]。说实话，TPU 我用的不多，但如果你做大规模训练，它确实香。

2.5 验证安装——跑个简单的例子

装完了，总得试试能不能用。我一般跑个矩阵乘法：

import jax.numpy as jnp
from jax import random

key = random.PRNGKey(0)
x = random.normal(key, (1000, 1000))
y = jnp.dot(x, x.T)
print(y.shape)

如果没报错，说明 JAX 能正常计算。你想想看，这一步其实挺重要的——很多问题都是在运行时才暴露出来。

2.6 配置开发环境——VS Code

我个人习惯用 VS Code，插件装好之后，写代码体验很不错。推荐几个插件：

Python：微软官方，必装
Jupyter：支持 notebook 交互
Pylance：代码补全和类型检查

配置好之后，按 Ctrl+Shift+P，选择 Python 解释器，指向你刚才创建的 conda 环境。这样就能在 VS Code 里直接跑 JAX 代码了。

2.7 配置开发环境——Jupyter Notebook

如果你喜欢交互式开发，Jupyter 是个好选择。安装：

pip install jupyter
jupyter notebook

然后在 notebook 里导入 JAX，跑几个 cell 试试。我经常在 Jupyter 里做原型验证，方便得很。

2.8 知识体系总览

下面这张图，帮你理清本章的核心逻辑：

2.9 常见问题与解决方案

问题	原因	解决方案
安装时报错 `No matching distribution found`	Python版本或CUDA版本不匹配	检查官方兼容性表，调整版本
GPU版本装上了，但 `jax.devices()` 只显示CPU	cuDNN未安装或版本不对	安装对应版本的cuDNN，或重新安装jaxlib
TPU版本报 `libtpu.so not found`	TPU驱动未正确安装	确认在TPU虚拟机上运行，或重新安装jax[tpu]
VS Code无法识别JAX	Python解释器未指向正确环境	按 `Ctrl+Shift+P` 选择正确的conda环境

💡 小技巧： 如果你在安装过程中遇到任何问题，先试试 pip install --upgrade jax jaxlib。很多时候，版本过旧是罪魁祸首。

2.10 总结

环境搭建这事儿，说白了就是三步：装JAX、验证、配工具。每一步都有坑，但只要你按部就班来，基本不会出大问题。我个人建议，先从CPU版本开始，跑通了再上GPU或TPU。这样即使出问题，也能快速定位。

📚 课程目录