一、JAX初探：为什么选择JAX？

大家好，我是你们这门课的老师。今天咱们聊聊JAX。

说实话，我第一次接触JAX是在2020年。当时我在做一个强化学习项目，PyTorch的训练速度让我有点头疼。同事推荐说「试试JAX吧」，我一开始是拒绝的——毕竟PyTorch用得好好的，干嘛要换？

结果一试，真香。

1.1 为什么选择JAX？

先说说JAX解决了什么问题。

你想想看，做深度学习最头疼的是什么？

• 训练太慢——GPU利用率上不去
• 代码调试麻烦——梯度爆炸了都不知道在哪一步出的问题
• 模型部署困难——训练环境和推理环境不一致

JAX就是冲着这些痛点来的。它不是一个完整的深度学习框架（像PyTorch或TensorFlow那样），而是一个高性能数值计算库。它的核心卖点就三个：

1. 自动微分——你写前向计算，它自动给你反向传播
2. JIT编译——把Python代码编译成高效的XLA代码，跑在GPU/TPU上
3. 函数式编程——没有副作用，代码可预测、可调试

1.2 JAX与NumPy的关系

JAX和NumPy的关系，说白了就是「长得像，但内核完全不同」。

你看这段代码：

import jax.numpy as jnp
import numpy as np

# NumPy写法
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
c = a + b  # 结果是 [5, 7, 9]

# JAX写法
a_j = jnp.array([1, 2, 3])
b_j = jnp.array([4, 5, 6])
c_j = a_j + b_j  # 结果也是 [5, 7, 9]

是不是几乎一模一样？

但区别在哪呢？

特性	NumPy	JAX
执行方式	即时执行（Eager）	可JIT编译
GPU支持	不支持	原生支持
自动微分	不支持	支持（grad函数）
数组不可变	可变	不可变
随机数生成	全局状态	显式状态管理

我在项目中遇到过一个问题：用NumPy写的预处理代码，迁移到JAX后速度反而慢了。后来发现是因为没有用JIT编译。加了@jit装饰器后，速度提升了20倍。

1.3 JAX的核心设计哲学

JAX的设计哲学，说白了就两句话：函数式编程和不可变性。

函数式编程

在JAX里，一切都是函数。没有类，没有对象，没有状态。

你写一个函数，JAX帮你做三件事：

• 自动微分：grad(f) 返回f的梯度函数
• JIT编译：jit(f) 返回编译后的高效版本
• 向量化：vmap(f) 自动处理批量数据

举个例子：

import jax
import jax.numpy as jnp

def f(x):
    return x ** 2 + 3 * x + 1

# 自动求导
grad_f = jax.grad(f)
print(grad_f(2.0))  # 输出 7.0 (因为导数 2x+3)

# JIT编译
fast_f = jax.jit(f)
print(fast_f(jnp.array([1.0, 2.0, 3.0])))

为什么会这样？因为函数式编程让JAX可以追踪计算图。你写的每个函数，JAX都能看到它的内部结构，然后做优化。

不可变性

这个点我刚开始很不习惯。在NumPy里，你可以直接修改数组：

# NumPy：可以原地修改
x = np.array([1, 2, 3])
x[0] = 100  # 没问题

但在JAX里，不行：

# JAX：会报错
x = jnp.array([1, 2, 3])
x[0] = 100  # TypeError: '' object does not support item assignment

你必须用函数式的方式：

x = jnp.array([1, 2, 3])
x_new = x.at[0].set(100)  # 返回新数组，原数组不变

我曾经在调试一个模型时，因为不小心修改了输入数组，导致梯度计算全错了。排查了整整一天才发现问题。从那以后，我彻底接受了「不可变性」这个设计。

1.4 JAX的安装与环境配置

安装JAX其实很简单。但有几个坑，我帮你踩过了。

基础安装

# CPU版本（最稳定）
pip install jax jaxlib

# GPU版本（需要CUDA）
pip install jax[cuda12] jaxlib

注意：GPU版本需要先装好CUDA和cuDNN。我建议用conda环境管理，避免依赖冲突。

验证安装

import jax
print(jax.__version__)  # 检查版本
print(jax.devices())    # 检查可用设备

# 如果看到 [CpuDevice(id=0)] 说明是CPU
# 如果看到 [GpuDevice(id=0)] 说明GPU可用

我的推荐配置

组件	推荐版本	备注
Python	3.10 - 3.12	3.9以下不推荐
JAX	0.4.20+	最新稳定版
CUDA	12.x	GPU版本必须
cuDNN	8.9+	与CUDA匹配

知识体系总览

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

嗯，以上就是第一章的全部内容。JAX的核心理念其实不复杂，但需要你动手去感受。下一章我们会深入JAX的自动微分机制，到时候见。

---