第2章环境搭建

第2章：环境搭建：安装JAX（CPU/GPU/TPU版本）、验证安装、JAX的依赖库

说实话，学任何框架的第一步，都是先把环境搭好。这一步要是卡住了，后面再好的内容也白搭。JAX 的安装其实不算复杂，但有几个坑我得提前跟你说清楚。

我个人习惯是，先搞清楚自己手头有什么硬件。是普通笔记本？还是有 GPU 的台式机？或者你运气好，能搞到 TPU？不同的硬件，安装命令差不少。

2.1 安装前的准备：Python 版本与虚拟环境

JAX 对 Python 版本有要求。我建议你用 Python 3.9 到 3.11 之间。3.12 虽然也能用，但有些依赖库还没完全跟上，容易出幺蛾子。

嗯，这里要注意：千万别直接在系统 Python 里装。你想想看，万一搞崩了系统环境，那可就麻烦了。我一般用 conda 或者 venv 创建虚拟环境。

我的习惯：用 conda 创建环境，名字就叫 jax_env。干净又卫生。

# 创建虚拟环境
conda create -n jax_env python=3.10

# 激活环境
conda activate jax_env

2.2 安装 CPU 版本的 JAX

如果你只是想在笔记本上先跑跑看，或者做点小实验，CPU 版本完全够用。安装命令非常简单：

pip install --upgrade pip
pip install --upgrade jax jaxlib

就这么两行。JAX 会自动帮你装好 CPU 版本的 jaxlib。我在项目中遇到过，有些同学用 pip install jax 之后发现 import 报错，其实就是 jaxlib 没装对。所以建议你两条命令一起跑。

2.3 安装 GPU 版本的 JAX

GPU 版本才是 JAX 真正发挥实力的地方。但这里有个大坑：CUDA 和 cuDNN 的版本必须匹配。

我曾经因为 CUDA 版本不对，折腾了一整个下午。后来学乖了，直接去 JAX 官方文档查对应关系。

JAX 版本	CUDA 版本	cuDNN 版本
0.4.20	11.8	8.6
0.4.23	12.1	8.9
0.4.28	12.2	8.9

确认好 CUDA 版本后，安装命令如下：

# 以 CUDA 12.1 为例
pip install --upgrade pip
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

避坑指南：不要用 pip install jax[cuda]，这个写法已经过时了。一定要指定 CUDA 版本号，比如 jax[cuda12_pip]。

2.4 安装 TPU 版本的 JAX

TPU 版本比较特殊。你需要在 Google Cloud 上创建 TPU 虚拟机，或者用 Colab 的 TPU 运行时。我个人觉得，初学者没必要一上来就搞 TPU，先把 CPU 和 GPU 玩明白再说。

如果你确实需要，命令如下：

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

2.5 验证安装是否成功

装完之后，别急着跑。先验证一下。我每次装完新环境，都会跑下面这段代码：

import jax
import jax.numpy as jnp

# 查看 JAX 版本
print(f"JAX version: {jax.__version__}")

# 查看后端设备
print(f"Devices: {jax.devices()}")

# 跑一个简单的矩阵乘法
a = jnp.ones((3, 4))
b = jnp.ones((4, 5))
c = jnp.dot(a, b)
print(f"Matrix multiplication result shape: {c.shape}")

如果输出类似下面这样，说明安装成功了：

JAX version: 0.4.28
Devices: [CpuDevice(id=0)]
Matrix multiplication result shape: (3, 5)

GPU 版本的话，jax.devices() 会显示 [GpuDevice(id=0)]。如果显示的是 CPU，那说明 GPU 没识别到。别慌，检查一下 CUDA 版本和 jaxlib 版本是否匹配。

2.6 JAX 的核心依赖库

光有 JAX 还不够。实际项目中，我们通常还会装几个配套库。说白了，JAX 是发动机，但这些库是方向盘和轮胎。

2.6.1 jaxlib

这是 JAX 的底层实现。它包含了 XLA 编译器和设备相关的代码。你装 JAX 的时候，它通常会自动装好。但如果遇到版本不匹配，可以单独指定版本安装。

2.6.2 optax

这是 JAX 的优化器库。我刚开始用 JAX 的时候，发现它居然没有内置的 Adam 优化器。后来才知道，官方把优化器都放到了 optax 里。

pip install optax

optax 提供了各种优化器，比如 Adam、SGD、AdamW 等。而且它的设计非常函数式——每个优化器都是一个函数，返回一个更新状态的函数。这个设计理念，说实话，比 PyTorch 的优化器要优雅得多。

2.6.3 flax

Flax 是 JAX 生态里最流行的神经网络库。它提供了 Layer、Module、训练循环等高层封装。

pip install flax

我在项目中遇到过，有人用 JAX 手写了一个 ResNet，写了 500 行代码。后来改用 Flax，100 行就搞定了。所以，别跟自己过不去，该用框架就用框架。

2.7 知识体系总览

下面这张图，是我整理的本章节核心逻辑。你可以把它当作一张地图，随时回来对照。

2.8 常见问题与避坑指南

问题1：安装后 import jax 报错 "No module named 'jaxlib'"

原因：jaxlib 没装成功。解决办法：手动指定版本安装。

pip install jaxlib==0.4.28

问题2：GPU 版本安装后，jax.devices() 只显示 CPU

原因：CUDA 版本不匹配，或者 jaxlib 装成了 CPU 版本。我曾经犯过这个错，后来发现是 pip 缓存搞的鬼。

解决办法：先清缓存，再指定 CUDA 版本安装。

pip cache purge
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

小技巧：如果你不确定自己的 CUDA 版本，可以用 nvidia-smi 查看。右上角会显示 CUDA Version。

好了，环境搭好了，依赖库也装齐了。接下来你就可以正式开始 JAX 的函数式编程之旅了。记住，环境问题是最容易解决的——无非就是版本匹配那点事。别怕，多试几次就熟了。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录