2、安装与第一行代码：JAX 的安装（CPU/GPU/TPU）、JAX 与 NumPy 的 API 对比、第一个 JAX 程序

好，咱们正式开始动手。这一章我带你走完 JAX 的安装流程，然后写第一行代码。别小看这一步，我见过不少人在安装环节就卡住了，尤其是 GPU 版本。

2.1 安装 JAX：CPU / GPU / TPU 版本

JAX 的安装其实比 PyTorch 要清爽不少。它没有那么多复杂的依赖，核心就一个包。但不同硬件版本，命令略有区别。

CPU 版本（最稳妥）

如果你只是想先体验一下，或者手头没有 NVIDIA 显卡，直接装 CPU 版本就行。我个人习惯先用 CPU 版本跑通逻辑，再切到 GPU 上做大规模训练。

pip install jax jaxlib

这条命令会安装纯 CPU 版本的 JAX。注意，jaxlib 是底层实现，必须一起装。

GPU 版本（CUDA 11 / 12）

如果你有 NVIDIA 显卡，想用 GPU 加速，那就得指定 CUDA 版本。这里有个坑——JAX 的 GPU 版本目前只支持 CUDA 11.8 或 12.x。我去年帮一个团队排查问题，发现他们装了 CUDA 11.4，结果 JAX 死活不认 GPU。折腾了半天，最后升级到 CUDA 11.8 才解决。

# CUDA 12
pip install --upgrade jax jaxlib==0.4.28+cuda12 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11.8
pip install --upgrade jax jaxlib==0.4.28+cuda11 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

TPU 版本（云上专用）

TPU 版本一般只在 Google Cloud 上使用。如果你没有 TPU 环境，可以跳过。但如果你在 Colab 上跑，JAX 已经预装了。

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

2.2 验证安装：你的第一个 JAX 程序

装完之后，咱们写个最简单的程序验证一下。打开 Python 交互环境，或者新建一个 .py 文件。

import jax
import jax.numpy as jnp

# 检查设备
print(jax.devices())

# 创建一个张量
x = jnp.array([1.0, 2.0, 3.0])
print(x)

如果输出类似 [TFRT_CPU_0] 或者 [CUDA_0]，说明安装成功。我第一次看到 CUDA_0 出现时，心里还是挺爽的——GPU 加速就这么简单。

2.3 JAX 与 NumPy 的 API 对比

JAX 的 API 设计几乎完全复刻了 NumPy。你想想看，这意味着什么？意味着你只要会 NumPy，就能直接上手 JAX。但有几个关键区别，我列个表格给你看。

功能	NumPy	JAX (jnp)	区别说明
数组创建	np.array([1,2,3])	jnp.array([1,2,3])	语法完全一致
随机数	np.random.randn(3)	jax.random.normal(key, (3,))	JAX 需要显式传入 key（随机种子）
矩阵乘法	np.dot(a, b)	jnp.dot(a, b)	语法一致，但 JAX 自动支持 GPU 加速
原地修改	a[0] = 1	不支持	JAX 数组不可变，必须用 index_update
梯度计算	无	jax.grad(f)	JAX 原生支持自动微分

说白了，JAX 的 jnp 就是 NumPy 的「升级版」——保留了同样的接口，但底层换成了 XLA 编译和自动微分。你写 NumPy 代码时，只要把 np 换成 jnp，大部分情况都能直接跑。但有两个地方要注意：

• 随机数需要 key：JAX 是函数式风格，没有全局随机状态。每次生成随机数都要传入一个 key。我刚开始也觉得麻烦，但后来发现这其实避免了「随机种子污染」的问题。
• 数组不可变：你不能写 x[0] = 1，会报错。必须用 jax.ops.index_update 或者干脆创建新数组。

2.4 第一个有意义的 JAX 程序：线性回归

光打印数组太无聊了。咱们写一个真正能体现 JAX 优势的程序——用梯度下降做线性回归。这段代码会用到 jax.grad，也就是自动微分。

import jax
import jax.numpy as jnp

# 生成模拟数据
key = jax.random.PRNGKey(42)
x = jax.random.normal(key, (100,))
true_w, true_b = 2.0, -1.0
y = true_w * x + true_b + 0.1 * jax.random.normal(key, (100,))

# 定义损失函数
def loss(w, b):
    pred = w * x + b
    return jnp.mean((pred - y) ** 2)

# 计算梯度
grad_loss = jax.grad(loss, argnums=(0, 1))

# 梯度下降
w, b = 0.0, 0.0
lr = 0.1
for i in range(100):
    dw, db = grad_loss(w, b)
    w -= lr * dw
    b -= lr * db
    if i % 20 == 0:
        print(f"Step {i}: w={w:.3f}, b={b:.3f}, loss={loss(w, b):.4f}")

print(f"Final: w={w:.3f}, b={b:.3f} (true: w=2.0, b=-1.0)")

看到没？jax.grad 自动帮我们算出了梯度。你不需要手写导数公式，也不需要反向传播的代码。这就是 JAX 的核心魅力——自动微分 + GPU 加速，两行代码搞定。

2.5 本章知识体系图

下面这张图帮你梳理了本章的核心逻辑：从安装到第一个程序，再到 API 对比，最后用自动微分收尾。

嗯，到这里你已经完成了 JAX 的安装和第一个程序。说实话，安装这一步虽然简单，但 GPU 版本确实容易踩坑。如果你遇到问题，记得检查 CUDA 版本和 jaxlib 的匹配关系。下一章咱们会深入自动微分机制，看看 jax.grad 背后到底发生了什么。

---