第一章：JAX基础回顾

1.1 JAX到底是什么？

说实话，我第一次接触JAX的时候，也以为它只是又一个深度学习框架。但用了一段时间后，我发现——这玩意儿跟PyTorch、TensorFlow的思路完全不一样。

JAX，全称是「Just After eXecution」。它本质上是一个可组合的函数变换库。什么意思呢？就是你写一个普通的Python函数，JAX能帮你自动求导、自动编译、自动并行化。嗯，听起来很玄乎，但用起来真的很爽。

我个人习惯把JAX看作「NumPy + 自动微分 + GPU加速 + 编译器」的合体。它不像PyTorch那样有复杂的动态图机制，而是更接近函数式编程的风格。说白了，你写一个函数，JAX帮你把它变成高性能的版本。

1.2 JAX与NumPy的关系

如果你用过NumPy，那JAX的上手成本几乎为零。JAX重新实现了NumPy的大部分API，就在jax.numpy里。你可以直接把它当成NumPy来用，只不过它跑在GPU上。

举个例子：

import jax.numpy as jnp
import numpy as np

# NumPy写法
a = np.array([1.0, 2.0, 3.0])
print(a.mean())  # 2.0

# JAX写法（几乎一样）
b = jnp.array([1.0, 2.0, 3.0])
print(b.mean())  # 2.0

我在项目中遇到过一个问题：刚开始用JAX时，我习惯性地调用了np.array，结果发现数据还在CPU上，根本没用到GPU。后来才意识到，必须用jnp.array才能享受加速。这个坑，我踩过。

对比项	NumPy	JAX
运行设备	CPU	CPU/GPU/TPU
自动微分	不支持	支持（grad）
JIT编译	不支持	支持（jit）
原地修改	支持	不支持（函数式）

1.3 JAX的核心概念

JAX的三大核心变换，你必须要掌握。它们分别是 jit、vmap 和 pmap。我一个个来讲。

1.3.1 jit：即时编译

jit 是JAX最常用的变换。它能把你的Python函数编译成高效的XLA代码。说白了，就是让你的代码跑得更快。

import jax
import jax.numpy as jnp

def slow_fn(x):
    for _ in range(1000):
        x = x * 2 + 1
    return x

# 不加jit，慢慢跑
# 加了jit，飞起来
fast_fn = jax.jit(slow_fn)

x = jnp.array([1.0, 2.0, 3.0])
result = fast_fn(x)
print(result)

我曾经在训练一个大型模型时，发现每次前向传播都要好几秒。加了jit之后，直接降到了毫秒级。你想想看，这差距有多大。

1.3.2 vmap：自动向量化

vmap 是JAX的另一个杀手锏。它能自动把你的函数变成批量处理的版本。不用手动写for循环，也不用自己处理batch维度。

import jax
import jax.numpy as jnp

def single_fn(x, y):
    return x * y + 1

# 手动批量处理
def batch_fn(X, Y):
    return jnp.array([single_fn(x, y) for x, y in zip(X, Y)])

# 用vmap自动向量化
batch_fn_vmap = jax.vmap(single_fn)

X = jnp.array([1.0, 2.0, 3.0])
Y = jnp.array([4.0, 5.0, 6.0])

print(batch_fn(X, Y))      # 手动版
print(batch_fn_vmap(X, Y)) # vmap版，结果一样

嗯，这里要注意：vmap 不只是语法糖。它底层做了优化，比手动for循环快得多。我在做强化学习时，经常用vmap来并行处理多个环境的状态。

1.3.3 pmap：多设备并行

pmap 是vmap的升级版。它把计算分布到多个设备（比如多个GPU）上。每个设备处理一部分数据，然后自动汇总结果。

import jax
import jax.numpy as jnp

# 假设你有4个GPU
def fn(x):
    return x ** 2

# 自动分布到4个GPU上
parallel_fn = jax.pmap(fn)

# 注意：输入的第一个维度会被切分到各个设备
x = jnp.array([[1.0], [2.0], [3.0], [4.0]])
result = parallel_fn(x)
print(result)

1.4 JAX的安装与配置

安装JAX其实很简单。但要注意，不同硬件平台的安装方式不一样。

1.4.1 CPU版本（最通用）

pip install jax jaxlib

这个版本所有电脑都能装。但说实话，没有GPU加速，JAX的优势发挥不出来。

1.4.2 GPU版本（推荐）

pip install jax[cuda12]  # 适用于CUDA 12.x
# 或者
pip install jax[cuda11]  # 适用于CUDA 11.x

我曾经在配置GPU版本时踩过一个坑：CUDA版本必须和jaxlib的版本严格匹配。否则会报错说找不到CUDA库。我的建议是，先查一下你的CUDA版本：

nvcc --version

然后去JAX官方文档找对应的安装命令。

1.4.3 TPU版本（谷歌云专用）

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

这个版本一般用不到。除非你在谷歌云上跑实验。

1.4.4 验证安装

装完之后，跑一下这段代码，看看能不能正常使用：

import jax
print(jax.devices())  # 应该能看到你的设备列表
print(jax.__version__)  # 版本号

如果能看到类似 [CudaDevice(id=0)] 的输出，那就说明安装成功了。

本章知识体系

下面这张图，帮你梳理了本章的核心内容：

好了，第一章的内容就到这里。JAX的基础概念其实不多，但每个都很重要。尤其是jit和vmap，在后面的梯度累积和混合精度训练中会频繁用到。建议你先把这些基础打牢，后面学起来会轻松很多。

---

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