第2章环境准备

第2章：环境准备：安装JAX（CPU/GPU版本）、验证安装、配置Google Colab环境、JAX与NumPy的第一次亲密接触

好，咱们正式开始动手了。

这一章，说白了就是帮你把JAX请进家门，再让它跟NumPy打个招呼。环境配置这事儿，看着琐碎，但要是没弄对，后面跑代码时会让你抓狂。我当年第一次装JAX的GPU版本，就因为CUDA版本没对齐，折腾了整整一个下午。嗯，咱们今天争取十分钟搞定。

2.1 安装JAX：CPU版本

如果你只是想先体验一下JAX，或者手头没有NVIDIA显卡，那装CPU版本就够了。我个人习惯在干净的虚拟环境里装，避免跟其他包打架。

# 用pip直接装
pip install --upgrade jax jaxlib

# 或者用conda（推荐，依赖管理更省心）
conda install -c conda-forge jax jaxlib

装完之后，你可以打开Python验证一下：

import jax
print(jax.__version__)
# 应该能看到类似 0.4.x 的版本号

如果没报错，恭喜你，CPU版本已经就绪。

小提示： 我建议你装完立刻跑一下 jax.numpy 的基本运算，比如 jax.numpy.array([1,2,3])，确保导入没问题。有时候pip会静默失败，但import才暴露问题。

2.2 安装JAX：GPU版本

GPU版本才是JAX真正发力的地方。但这里有个坑——CUDA和cuDNN的版本必须跟JAX匹配。我曾经因为偷懒没查版本表，结果跑模型时一直报“CUBLAS_STATUS_NOT_INITIALIZED”，排查了半天才发现是CUDA 11.8跟JAX 0.4.10不兼容。

官方推荐的做法是：

# 先确认你的CUDA版本（比如11.8或12.x）
nvcc --version

# 然后安装对应版本的jaxlib
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# 如果你用的是CUDA 12，把上面cuda11改成cuda12

装完后，用下面这段代码验证GPU是否被识别：

import jax
print(jax.devices())
# 如果输出 [cuda(id=0)] 之类的，说明GPU版本装好了

注意： 千万别用 pip install jax[cuda] 这种写法，它已经过时了。现在必须指定CUDA大版本号。另外，Windows用户建议直接用WSL2，原生Windows支持还在实验阶段。

2.3 配置Google Colab环境

如果你不想折腾本地环境，或者手头没有GPU，Colab是最省事的选择。Colab默认已经装了JAX，但版本可能比较老。我建议你每次新建笔记本时，先升级一下：

!pip install --upgrade jax jaxlib

然后记得把运行时类型改成GPU：

点击菜单栏「运行时」→「更改运行时类型」
在「硬件加速器」里选「GPU」
保存后重启运行时

验证一下：

import jax
print(jax.devices())
# 应该能看到 [cuda(id=0)]

个人经验： Colab的免费GPU是T4，虽然不算顶级，但跑JAX的demo和中小型模型完全够用。我经常在Colab上做原型验证，确认没问题了再丢到本地服务器上跑大规模实验。

2.4 JAX与NumPy的第一次亲密接触

环境准备好了，咱们来写第一段JAX代码。你会发现，JAX的API跟NumPy几乎一模一样——说白了，它就是NumPy的“加速版”。

import jax.numpy as jnp
import numpy as np

# 创建数组
a_np = np.array([1.0, 2.0, 3.0])
a_jax = jnp.array([1.0, 2.0, 3.0])

print(a_np)
print(a_jax)
# 输出看起来几乎一样

但别被表象骗了。JAX的数组是不可变的。你想想看，NumPy里你可以直接 a[0] = 100 原地修改，但在JAX里这么干会报错：

# 这行会报错！
a_jax[0] = 100.0
# TypeError: '' object does not support item assignment

为什么会这样？因为JAX追求函数式纯性——没有副作用，才能做自动微分和JIT编译。如果你想“修改”数组，得用 at 方法：

a_new = a_jax.at[0].set(100.0)
print(a_new)  # [100., 2., 3.]
print(a_jax)  # 原数组不变 [1., 2., 3.]

嗯，这里要注意：at 返回的是新数组，原数组纹丝不动。刚开始用可能会觉得别扭，但习惯之后你会发现，这种设计让代码的副作用一目了然，调试起来特别爽。

咱们再跑个矩阵乘法，感受一下性能差异：

import time

# 生成大矩阵
size = 5000
x_np = np.random.randn(size, size)
y_np = np.random.randn(size, size)

x_jax = jnp.array(x_np)
y_jax = jnp.array(y_np)

# NumPy 耗时
start = time.time()
z_np = np.dot(x_np, y_np)
print(f"NumPy: {time.time() - start:.3f} 秒")

# JAX 耗时（第一次运行会包含编译时间）
start = time.time()
z_jax = jnp.dot(x_jax, y_jax)
print(f"JAX (首次): {time.time() - start:.3f} 秒")

# 第二次运行，JIT编译缓存生效
start = time.time()
z_jax = jnp.dot(x_jax, y_jax)
print(f"JAX (第二次): {time.time() - start:.3f} 秒")

你会发现，JAX第一次跑时可能比NumPy还慢——因为它在背后做JIT编译。但第二次开始，速度就起飞了。我当年第一次看到这个对比时，心里就一个想法：这玩意儿，真香。

  核心要点： JAX的API设计让你可以用NumPy的思维写代码，但底层机制完全不同。记住三点：数组不可变、函数式风格、JIT编译加速。

2.5 本章知识体系

下面这张图帮你梳理了本章的核心逻辑：从环境搭建到第一个JAX程序，每一步都有它的意义。

这张图把本章的脉络串起来了：先装CPU或GPU版本，再验证设备是否识别，最后用JAX写第一段代码，感受它跟NumPy的异同。每一步都是后面深入学习的基础。

避坑指南： 我曾经在GPU版本上花了两小时，最后发现只是CUDA路径没加到环境变量里。如果你遇到 Could not find cuda 之类的错误，先检查 echo $LD_LIBRARY_PATH 有没有包含CUDA的lib目录。另外，JAX的GPU版本对cuDNN版本也很敏感，建议直接用conda安装，它会自动帮你处理依赖。

好了，环境已经搭好，JAX也跑起来了。下一章咱们会深入JAX的核心——自动微分，那才是JAX真正让人拍大腿的地方。

📚 课程目录