第2章 JAX基础速通

第2章：JAX基础速通

好，咱们直接进入正题。JAX 这东西，说白了就是 Google 搞出来的一个“可微分的 NumPy”。我刚开始接触的时候，心里也犯嘀咕：这不就是另一个深度学习框架吗？但用了一段时间后，我发现它跟 PyTorch 和 TensorFlow 的思路完全不一样。它更像是一个“函数变换器”——你给它一个普通的 Python 函数，它能帮你变出梯度、变出向量化版本、甚至变出能在 GPU 上跑的编译版本。

这一章，咱们就快速过一遍 JAX 最核心的四个能力：NumPy 兼容 API、JIT 编译、自动微分（grad）和向量化映射（vmap）。这些都是你后面做贝叶斯建模的“基本功”。

2.1 NumPy 兼容 API：几乎零成本迁移

JAX 的 API 设计，说白了就是“抄”NumPy 的。你想想看，NumPy 的接口大家用了多少年了？JAX 直接拿来用，你几乎不需要学新东西。我个人的习惯是，把 NumPy 的 import 语句从 import numpy as np 改成 import jax.numpy as jnp，然后代码基本就能跑。

核心区别：JAX 的数组默认是不可变的（immutable）。这意味着你不能像 NumPy 那样直接原地修改数组元素。比如 arr[0] = 1 会报错。你得用 arr.at[0].set(1) 这种函数式的方式。

import jax.numpy as jnp
import numpy as np

# 创建数组，跟 NumPy 一模一样
a = jnp.array([1.0, 2.0, 3.0])
b = jnp.ones((3, 3))

# 数学运算，完全兼容
c = jnp.dot(a, a)
print(c)  # 输出 14.0

# 注意：不能原地修改！
# a[0] = 10.0  # 这行会报错
a = a.at[0].set(10.0)  # 正确做法
print(a)  # 输出 [10. 2. 3.]

嗯，这里要注意：JAX 的数组默认在 CPU 上。如果你想用 GPU，JAX 会自动检测。你不需要写 .cuda() 之类的代码。我刚开始用的时候，还傻乎乎地到处找 .to(device) 的写法，后来发现根本不需要。

2.2 JIT 编译：让代码飞起来

JIT 编译，全称是“Just-In-Time”编译。它的作用是把你的 Python 函数编译成 XLA（Accelerated Linear Algebra）的底层表示，然后在 GPU 或 TPU 上高效执行。说白了，就是让你的代码跑得更快。

用法超级简单：用 @jit 装饰器装饰你的函数就行。我曾在项目中遇到过一个大循环，每次迭代都要做矩阵乘法。用 JIT 之后，速度提升了将近 20 倍。你想想看，这感觉有多爽？

from jax import jit
import jax.numpy as jnp

# 定义一个简单的函数
def slow_function(x, y):
    for _ in range(1000):
        x = jnp.dot(x, y)
    return x

# 用 JIT 编译
@jit
def fast_function(x, y):
    for _ in range(1000):
        x = jnp.dot(x, y)
    return x

# 测试
x = jnp.ones((100, 100))
y = jnp.ones((100, 100))

# 第一次调用会触发编译，稍慢
result = fast_function(x, y)
# 后续调用飞快
result = fast_function(x, y)

避坑指南：JIT 编译对 Python 的控制流（如 if、for）有限制。我曾经写过一个带条件分支的函数，结果 JIT 编译报错了。解决方案是使用 jax.lax.cond 或 jax.lax.fori_loop 等函数式控制流。

2.3 自动微分（grad）：贝叶斯建模的基石

自动微分是 JAX 最强大的功能之一。你只需要定义一个函数，然后用 grad 函数就能得到它的梯度。这对于贝叶斯建模来说至关重要——因为我们要计算后验分布的梯度，才能用梯度下降法或 HMC 采样。

我个人觉得，JAX 的自动微分比 PyTorch 的 backward() 更直观。它不修改你的函数，而是返回一个新的函数，这个新函数计算原函数的梯度。你想想看，这多符合函数式编程的思想？

from jax import grad
import jax.numpy as jnp

# 定义一个简单的函数
def f(x):
    return x**2 + 3*x + 1

# 计算梯度
grad_f = grad(f)

# 在 x=2.0 处求梯度
print(grad_f(2.0))  # 输出 7.0（因为导数是 2x+3）

# 高阶梯度？没问题！
grad_grad_f = grad(grad(f))
print(grad_grad_f(2.0))  # 输出 2.0（二阶导数）

小技巧：如果你需要计算多个参数的梯度，可以用 grad(f, argnums=(0, 1)) 指定对哪些参数求导。我在做贝叶斯线性回归时，经常需要对权重和偏置同时求梯度，这个参数就派上用场了。

2.4 向量化映射（vmap）：告别手动批处理

vmap 是 JAX 里一个“神奇”的功能。它的作用是把一个处理单个样本的函数，自动变成能处理批量样本的函数。你不需要手动写 for 循环，也不需要调整维度。vmap 会自动帮你“向量化”。

我举个例子。假设你有一个函数，计算两个向量的点积。如果你有一批向量，想计算两两之间的点积，传统做法是写个循环。但用 vmap，一行代码搞定。

from jax import vmap
import jax.numpy as jnp

# 定义一个处理单个样本的函数
def dot_product(x, y):
    return jnp.dot(x, y)

# 创建一批数据
batch_x = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
batch_y = jnp.array([[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]])

# 用 vmap 向量化
batched_dot = vmap(dot_product)

# 计算批量点积
results = batched_dot(batch_x, batch_y)
print(results)  # 输出 [23. 67. 131.]

嗯，这里要注意：vmap 不仅适用于简单的函数，也适用于复杂的神经网络层。我在做贝叶斯神经网络时，经常用 vmap 来批量计算每个数据点的对数似然，然后求和。这比手动写循环快多了。

2.5 知识体系总览

为了让你更直观地理解 JAX 的核心逻辑，我画了一张图。这张图展示了 JAX 的四大核心能力如何协同工作：

从这张图你可以看到，JAX 的四大核心能力是相互独立的，但又可以组合使用。比如，你可以先定义一个函数，用 grad 求梯度，然后用 jit 编译加速，最后用 vmap 批量处理数据。这种组合能力，正是 JAX 在贝叶斯建模中如此强大的原因。

2.6 实战小练习：组合使用

光说不练假把式。咱们来个小练习，把今天学的四个能力组合起来用。假设我们要计算一个简单函数的梯度，并对一批数据做向量化计算：

from jax import grad, jit, vmap
import jax.numpy as jnp

# 定义函数：f(x, w) = w * x^2
def f(x, w):
    return w * x**2

# 对 w 求梯度
grad_f = grad(f, argnums=1)

# 用 JIT 编译
grad_f_jit = jit(grad_f)

# 创建一批数据
xs = jnp.array([1.0, 2.0, 3.0])
ws = jnp.array([0.5, 1.0, 1.5])

# 用 vmap 批量计算梯度
batched_grad = vmap(grad_f_jit, in_axes=(0, 0))

# 计算
results = batched_grad(xs, ws)
print(results)  # 输出每个样本的梯度

你看，代码很简洁，但背后做了很多事情：自动微分、JIT 编译、向量化映射。这就是 JAX 的魅力所在。

我的建议：刚开始学 JAX 时，不要急着用高级功能。先把 NumPy 兼容 API 用熟，然后逐步引入 JIT、grad 和 vmap。我在教团队新人时，都是让他们先写纯 NumPy 代码，再一步步“JAX 化”。这样不容易出错。

好了，这一章的内容就到这里。你掌握了 JAX 的四大核心能力，后面做贝叶斯建模就会轻松很多。记住，JAX 的核心思想是“函数变换”——你给它一个函数，它帮你变出各种有用的版本。下一章，咱们会深入探讨 JAX 的概率编程库，看看如何用 JAX 做贝叶斯推断。

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

📚 课程目录