3、JAX核心概念：函数式编程、不可变状态、纯函数、JIT编译原理

好，咱们直接进入正题。JAX 这套框架，说实话，刚接触的时候挺别扭的。你想想看，我们平时写 PyTorch 或者 TensorFlow，习惯了那种“先定义模型、再 forward、再 backward”的套路。但 JAX 不一样，它逼着你换一种思维方式——函数式编程。

我个人习惯把 JAX 的核心思想总结成四个词：纯函数、不可变状态、函数式变换、JIT 编译。这四个词搞明白了，JAX 的底层逻辑你就拿捏了八成。

3.1 纯函数：没有副作用，才是好函数

什么叫纯函数？说白了就是：同样的输入，永远得到同样的输出，而且不修改任何外部状态。

我在项目中遇到过这样一个坑：写了一个函数，里面偷偷改了一个全局变量。在 Python 里这很正常，但在 JAX 里，这会导致 JIT 编译出来的结果完全不可控。为什么呢？因为 JIT 会缓存函数的编译结果，如果你的函数有副作用，缓存的结果就可能跟实际执行不一致。

# 不纯的函数 —— 在 JAX 里会出问题
global_var = 0
def impure_add(x):
    global global_var
    global_var += 1
    return x + global_var

# 纯函数 —— JAX 喜欢这样
def pure_add(x, y):
    return x + y

嗯，这里要注意：JAX 的 jit 装饰器会假设你传入的函数是纯函数。如果你在里面搞了副作用，JIT 编译出来的结果可能会让你怀疑人生。

3.2 不可变状态：数组一旦创建，就不能改

这个更狠。在 NumPy 里，你可以写 arr[0] = 1 直接修改数组。但在 JAX 里，数组是不可变的。

你可能会问：“那我怎么更新参数？”答案是：创建新的数组。

import jax.numpy as jnp

x = jnp.array([1, 2, 3])
# x[0] = 10  # 这行会报错！JAX 数组不可变

# 正确做法：创建新数组
y = x.at[0].set(10)  # y = [10, 2, 3]，x 保持不变

我曾经在调优一个强化学习模型时，因为没注意这个特性，在循环里反复创建新数组，导致内存暴涨。后来才意识到：不可变状态虽然安全，但频繁创建新对象会有性能开销。解决方案是用 jit 把整个更新逻辑编译成一个函数，这样 JAX 会在底层做内存复用优化。

3.3 函数式变换：JAX 的杀手锏

JAX 最让我惊艳的地方，就是它的函数式变换。说白了，就是把函数当成数据，可以随意变换。

常见的变换有四种：

变换	作用	类比
jit	即时编译，加速执行	把 Python 函数编译成 XLA 计算图
grad	自动求导	对函数求梯度，返回一个新函数
vmap	自动向量化	把处理单个样本的函数，变成处理批量的
pmap	自动并行化	把函数分布到多个设备上执行

你想想看，这些变换可以组合使用。比如：

import jax
import jax.numpy as jnp

def loss_fn(params, x, y):
    pred = jnp.dot(params, x)
    return (pred - y) ** 2

# 先求梯度，再 JIT 编译，再向量化
grad_fn = jax.grad(loss_fn)
jit_grad_fn = jax.jit(grad_fn)
batched_grad_fn = jax.vmap(jit_grad_fn, in_axes=(None, 0, 0))

# 现在 batched_grad_fn 可以一次性处理整个 batch 的数据

我个人习惯把这种组合叫做“变换流水线”。你可以在上面叠加任意多的变换，JAX 会帮你把它们融合成一个高效的底层实现。

3.4 JIT 编译原理：Python 代码如何变成 GPU 指令？

好，这是本章的重头戏。JIT 编译到底是怎么工作的？

简单来说，JAX 的 jit 做了三件事：

1. 追踪（Tracing）：用虚拟数据跑一遍你的函数，记录下所有操作，构建一个计算图。
2. 优化（Optimization）：XLA 编译器对这个计算图做各种优化，比如算子融合、常量折叠、内存规划。
3. 代码生成（Code Generation）：生成针对目标设备（CPU/GPU/TPU）的机器码。

嗯，这里有个关键点：JIT 编译发生在第一次调用时。第一次调用会花点时间做编译，后续调用就直接用缓存的结果，速度飞快。

来看一个具体的例子，理解 JIT 的“追踪”过程：

import jax
import jax.numpy as jnp

@jax.jit
def f(x, y):
    z = x + y
    w = jnp.sin(z)
    return w

# 第一次调用：触发编译
result = f(jnp.array([1.0]), jnp.array([2.0]))

# 第二次调用：直接用编译好的代码
result = f(jnp.array([3.0]), jnp.array([4.0]))

在第一次调用时，JAX 会创建两个“追踪器”（tracer）对象来代替真实的数组。这些追踪器会记录下所有操作：先做加法，再做 sin。然后 JAX 把这个操作序列发给 XLA，XLA 生成高效的 GPU 代码。

你可能会问：“那如果我的函数里有条件分支怎么办？”好问题。JAX 的追踪器会同时追踪所有分支，然后生成一个包含条件判断的计算图。但这也意味着：你的函数不能依赖 Python 的控制流（比如 if x > 0 这种），因为追踪器不知道 x 的具体值。解决方案是用 jax.lax.cond 或者 jax.lax.while_loop 这类函数式控制流。

3.5 知识体系总览

下面这张图，是我自己总结的 JAX 核心概念关系图。你看一眼，就能明白这些概念是怎么串起来的。

这张图把四个核心概念的关系讲清楚了。你看，纯函数和不可变状态是 JAX 的“哲学基础”，函数式变换是“工具箱”，JIT 编译是“性能引擎”。四者缺一不可。

3.6 实战中的感悟

最后聊点实际的。我刚开始用 JAX 的时候，最大的障碍就是思维转换。写惯了 PyTorch 的 nn.Module，突然要改成纯函数，真的很不习惯。

但后来我发现一个规律：把状态显式地传进传出函数，反而让代码更容易调试。因为每个函数都是自包含的，你不需要追踪全局变量的变化轨迹。这在分布式训练和模型并行时尤其有用——每个设备拿到的都是完整的状态副本，不会出现“参数不同步”的 bug。

嗯，如果你是从 PyTorch 转过来的，我建议你先从小函数开始练手。比如写一个纯函数的线性层，然后用 jit 编译，再用 grad 求导。等你习惯了这种“函数式”的节奏，再去看完整的训练循环，就会豁然开朗。

---

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