真实程序都有状态，例如训练步骤之间不断变化的模型参数、优化器动量、运行统计量和计数器。JAX 提供两种表达有状态计算的方法：

1. 让状态穿过纯函数：状态作为参数进入函数，又作为返回值离开。这是经典方式，也是几乎整个 JAX 生态的标准做法。
2. Ref：显式的可变数组。当逐层传递状态太麻烦，或者确实需要原地写内存时使用。

1. 先弄明白“状态”是什么

可以先把状态理解为：

一次函数调用结束后，需要保存下来，并且会影响下一次调用的信息。

计数器当前数到几、模型目前的参数、优化器记录的动量、训练到了第几步，都是状态。

1.1 没有状态的函数

def add_one(x):
    return x + 1

只要输入相同，输出就相同：

add_one(0)  # 1
add_one(0)  # 1
add_one(0)  # 1

函数不会记住自己之前被调用过几次，每次计算只取决于当前输入。这是一个无状态的纯函数。

1.2 隐藏在对象里的状态

普通 Python 计数器会把当前数字保存在对象属性里：

class Counter:
    def __init__(self):
        self.n = 0

    def count(self):
        self.n += 1
        return self.n

counter = Counter()
counter.count()  # 1
counter.count()  # 2
counter.count()  # 3

虽然每次调用 count() 都没有传入参数，结果却不同，因为 self.n 保存了上一次调用留下的信息：

调用前状态 0 → 输出 1 → 保存状态 1
调用前状态 1 → 输出 2 → 保存状态 2
调用前状态 2 → 输出 3 → 保存状态 3

这个程序不是没有状态，而是把状态隐藏在 Python 对象内部。只看 count() -> int 的函数接口，看不出它还读取和修改了 self.n。

1.3 把状态写成输入和输出

JAX 更喜欢把同一件事写成：

def count(state):
    new_state = state + 1
    output = new_state
    return output, new_state

调用者负责保存函数返回的新状态：

state = 0

value, state = count(state)  # value=1, state=1
value, state = count(state)  # value=2, state=2
value, state = count(state)  # value=3, state=3

现在完整的数据流出现在函数接口中：

旧状态 0 → count → 输出 1，新状态 1
旧状态 1 → count → 输出 2，新状态 2
旧状态 2 → count → 输出 3，新状态 3

函数没有偷偷修改外部对象，只根据输入计算并返回结果，因此 JAX 更容易对它进行 tracing、JIT 编译、求导和批处理。

1.4 “加状态”和“没加状态”的区别

严格来说，JAX 版本没有凭空增加状态。两个计数器本来都需要记住当前数字，区别只是状态放在哪里：

写法	状态在哪里	函数接口能否看到	JAX 是否容易追踪
self.n += 1	Python 对象内部	否	否
output, new_state = count(state)	函数参数和返回值	是	是
jax.new_ref(...)	JAX 管理的可变内存	是	可以，但有限制

最核心的区别是：

隐藏状态：函数偷偷读取或修改外面的数据
显式状态：旧状态作为输入，新状态作为输出

1.5 模型参数为什么也是状态

模型训练本质上也是一个状态不断更新的过程：

def train_step(params, batch):
    grads = compute_gradients(params, batch)
    new_params = update(params, grads)
    return new_params

params 是状态，因为它由之前的训练产生，会影响这一次模型的输出，并在本次训练后变成下一步参数：

参数₀ → 训练第 1 步 → 参数₁
参数₁ → 训练第 2 步 → 参数₂
参数₂ → 训练第 3 步 → 参数₃

如果每一步都从最初的 params_0 开始，就会丢弃上一步学到的结果：

# 错误：三次训练都从同一份旧参数开始。
new_params = train_step(params_0, batch_1)
new_params = train_step(params_0, batch_2)
new_params = train_step(params_0, batch_3)

正确做法是把新参数继续传给下一步：

params = params_0
params = train_step(params, batch_1)
params = train_step(params, batch_2)
params = train_step(params, batch_3)

使用 momentum 或 AdamW 时，除了参数还要保存历史梯度信息。比较完整的训练状态可能是：

state = {
    "params": params,        # 模型当前学到了什么
    "opt_state": opt_state,  # 优化器保存的历史信息
    "step": step,            # 当前训练步，影响学习率调度
    "random_key": key,       # 下一次 Dropout 或采样使用
}

因此，JAX 训练代码经常写成：

new_state, metrics = train_step(state, batch)

它表示：使用当前全部状态和一批数据完成一次训练，返回下一步需要的全部新状态，以及本次的 loss 等指标。

状态不是 JAX 额外添加的东西，而是程序本来就需要记住的信息。JAX 只是要求尽量把这些信息从对象内部拿出来，明确写成函数的输入和输出。

2. 不可追踪状态的问题

从一个普通 Python 计数器开始：

import jax
import jax.numpy as jnp

class Counter:
    """一个简单计数器。"""

    def __init__(self):
        self.n = 0

    def count(self) -> int:
        """计数加一，并返回新值。"""
        self.n += 1
        return self.n

counter = Counter()

for _ in range(3):
    print(counter.count())

输出是 1、2、3。属性 n 在调用之间保存状态，self.n += 1 通过副作用更新它。普通 Python 中没有问题，但放到 JAX 变换下就会失败。

用 jaxpr 查看 tracing 记录到的程序：

counter = Counter()
jax.jit(counter.count).trace().jaxpr

{ lambda ; . let  in (1:i32[],) }

这段被记录的程序没有输入、不执行任何操作，而且永远返回 1。self.n += 1 只在 tracing 时执行了一次；产生的 self.n 被捕获为常量，而状态更新没有进入 jaxpr。

3. 让状态穿过纯函数

解决办法是把状态变成函数接口的一部分：接收当前状态，同时返回输出和更新后的状态。

CounterState = int

class CounterV2:

    def count(self, n: CounterState) -> tuple[int, CounterState]:
        # 分开写是为了区分“对外输出”和“下一步状态”。
        return n + 1, n + 1

    def reset(self) -> CounterState:
        return 0

counter = CounterV2()
state = counter.reset()

for _ in range(3):
    value, state = counter.count(state)
    print(value)

调用者现在显式维护状态。作为交换，count 成为纯函数，容易被 tracing：

jax.jit(counter.count).trace(0).jaxpr

{ lambda ; a:i32[]. let
    b:i32[] = add a 1:i32[]
    c:i32[] = add a 1:i32[]
  in (b, c) }

状态清楚地流过被记录的程序：作为参数进入，作为结果离开。没有隐藏或固化的值，因此 JAX 变换能正确处理它。

一般而言，需要把有状态接口：

class StatefulClass:
    state: State

    def stateful_method(*args, **kwargs) -> Output:
        ...

改写为：

def stateless_method(
    state: State, *args, **kwargs
) -> tuple[Output, State]:
    ...

这是常见的函数式编程模式，也是绝大多数 JAX 程序表达状态的方式。

4. 完整示例：带动量的训练状态

训练循环是最常见的例子：参数和优化器状态逐步穿过纯 update 函数。

def update(params, opt_state, batch):
    """一个训练步骤：消费旧状态，产生新状态。"""
    grads = jax.grad(loss_fn)(params, batch)
    new_params, new_opt_state = optimizer_step(
        params, grads, opt_state
    )
    return new_params, new_opt_state

def loss_fn(params, batch):
    x, y = batch
    return jnp.mean((params["w"] * x + params["b"] - y) ** 2)

def optimizer_step(params, grads, opt_state, lr=0.1, decay=0.9):
    # 带动量的梯度下降：momentum 就是优化器状态。
    new_opt_state = jax.tree.map(
        lambda m, g: decay * m + g, opt_state, grads
    )
    new_params = jax.tree.map(
        lambda p, m: p - lr * m, params, new_opt_state
    )
    return new_params, new_opt_state

params = {"w": jnp.float32(1.0), "b": jnp.float32(0.0)}
opt_state = jax.tree.map(jnp.zeros_like, params)
batch = (
    jnp.array([1.0, 2.0, 3.0]),
    jnp.array([3.0, 5.0, 7.0]),
)

for step in range(100):
    params, opt_state = update(params, opt_state, batch)

print(jax.tree.map(lambda x: round(float(x), 2), params))
# {'b': 1.0, 'w': 2.0}，拟合出 y = 2x + 1

这里的状态是 PyTree，所以相同模式可以从一个标量计数器扩展到整个模型。Optax 等优化器库也围绕如下接口设计：

update(grads, opt_state, ...) -> (updates, new_opt_state)

将状态作为值传递还有一个好处：每份状态都是不可变快照，JAX 变换可以作用于整个循环。我们可以对更新步骤求导，或用 vmap 同时运行许多独立训练，而不必担心多个引用别名指向同一份可变数据。

5. Ref：可变数组

显式传递值是标准方法，但有时会很麻烦。例如调用栈深处的函数想更新归一化统计量或记录指标，沿途每个函数都要接收并返回这份状态。

JAX 因此提供 Ref：可以原地读写的可变数组。使用 jax.new_ref 创建：

x_ref = jax.new_ref(jnp.zeros(3))

@jax.jit
def f():
    x_ref[1] += 1.0

print(x_ref)  # Ref([0., 0., 0.])
f()
f()
print(x_ref)  # Ref([0., 2., 0.])

x_ref[...] 把完整值读成普通 Array，x_ref[...] = A 写入完整数组：

def g(x):
    x_ref = jax.new_ref(0.0)
    x_ref[...] = jnp.sin(x)
    return x_ref[...]

print(jax.grad(g)(1.0))  # 0.5403023

Ref 与 Array 是不同类型。索引读写几乎是 Ref 唯一允许的直接操作；不能在需要 Array 的地方直接传 Ref：

x_ref = jax.new_ref(1.0)
# jnp.sin(x_ref)       # 错误：不能直接对 Ref 做数学运算
jnp.sin(x_ref[...])    # 正确：先读取 Array 值

6. Ref 的索引和 freeze

读写支持 NumPy 索引表达式：

x_ref = jax.new_ref(jnp.arange(12.0).reshape(3, 4))

row = x_ref[0]                 # 整数索引
x_ref[1] = row
val = x_ref[1, 2]              # 元组索引
x_ref[2, 3] = val
col = x_ref[:, 1]              # 切片
x_ref[0, :3] = col

vals = x_ref[
    jnp.array([0, 0, 1]),
    jnp.array([1, 2, 3]),
]
x_ref[
    jnp.array([1, 2, 1]),
    jnp.array([0, 0, 1]),
] = vals

索引大体遵循 NumPy，但已提前知道的越界索引会报错。这与普通 Array 的越界读取截断、越界写入丢弃不同。

完成修改后，jax.freeze 会使 Ref 失效——之后再次访问会报错——并把最终值返回为普通不可变 Array：

final = jax.freeze(x_ref)
print(final)

[[ 1.  1.  9.  3.]
 [ 1.  3.  2.  3.]
 [ 9.  9. 10.  2.]]

7. Ref 与纯函数

函数纯度有三个主要好处：使用者更容易推理程序和变换；编译器更容易优化、并行和扩展；JAX 更容易 tracing。

JAX 会拦截 Ref 操作，因此 tracing 本身没有问题。Ref 会在一定程度上限制编译器变换代码，但限制程度与显式传递状态相当。需要学习的重点是 Ref 与各种变换如何交互。

如果函数只在内部创建和使用 Ref，从调用者角度看它仍然是纯函数：

def normalize(x):
    acc = jax.new_ref(0.0)
    acc[...] = jnp.sum(x)
    return x / acc[...]

result = jax.grad(lambda x: normalize(x).sum())(
    jnp.arange(1.0, 4.0)
)
print(result)  # [0., 0., 0.]

函数只有在接收 Ref 作为输入时才是不纯的：Ref 可以是显式参数，也可以通过闭包捕获。纯度取决于调用者能否观察到状态变化。

不纯函数仍支持许多变换。例如可以让 vmap 沿 Ref 的不同 batch 元素分别写入：

def scale_into(x, out_ref):
    out_ref[...] = 2.0 * x

xs = jnp.arange(3.0)
out_ref = jax.new_ref(jnp.zeros(3))
jax.vmap(scale_into)(xs, out_ref)
print(out_ref[...])  # [0., 2., 4.]

但不能 vmap 一个通过闭包捕获 Ref 的函数，因为每个 batch 成员都会写同一个共享位置，最终值存在歧义：

r = jax.new_ref(0.0)

def write_shared(x):
    r[...] = x

# 会报错：每个 batch 成员都试图写同一个 Ref。
# jax.vmap(write_shared)(jnp.arange(3.0))

8. Ref 的限制

• 不能从 JIT 函数或 jax.lax.scan 等控制流 body 返回 Ref；
• 不能把同一个 Ref 两次传给一个 JIT 函数，也不能既传入某个 Ref、又让函数从外层闭包捕获它；
• 只能在创建 Ref 的作用域中 freeze；
• 不支持 Ref 指向 Ref；
• 不能对闭包捕获 Ref 的函数使用 vmap 或 shard_map。

这些规则用于排除别名和其他语义不明确的情况。违反规则时，错误信息会说明原因；部分限制将来可能放宽。

Ref 也是性能工具：写 Ref 会原地更新 buffer，而不是分配新数组。详细性能行为属于 JIT 章节。Ref 还可用于从反向传播传出值、跨 microbatch 原地累积梯度，以及对 Ref 参数求导，详见进阶的 Autodiff with refs。

9. 语言模型训练状态

官方示例中的 params 与 opt_state 可以直接扩展成语言模型状态：

state = {
    "params": params,
    "opt_state": opt_state,
    "step": jnp.array(0, dtype=jnp.int32),
    "key": jax.random.key(0),
}

语言模型 logits 为 [B,T,V]，targets 为 [B,T]：

def cross_entropy(logits, targets):
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    target_log_probs = jnp.take_along_axis(
        log_probs, targets[..., None], axis=-1
    )[..., 0]
    return -jnp.mean(target_log_probs)

完整训练步骤仍遵循“旧状态进入、新状态离开”：

def loss_fn(params, batch, key):
    logits = model(params, batch["inputs"], key, training=True)
    return cross_entropy(logits, batch["targets"])

def train_step(state, batch, learning_rate):
    next_key, dropout_key = jax.random.split(state["key"])
    loss, grads = jax.value_and_grad(loss_fn)(
        state["params"], batch, dropout_key
    )
    new_params = jax.tree.map(
        lambda p, g: p - learning_rate * g,
        state["params"], grads,
    )
    new_state = {
        **state,
        "params": new_params,
        "step": state["step"] + 1,
        "key": next_key,
    }
    return new_state, {"loss": loss}

train_step_jit = jax.jit(train_step)

这里保留用 SGD，实际模型通常用 Optax 提供 AdamW、学习率调度和梯度裁剪；它仍遵循官方介绍的状态传递模式。checkpoint 应同时保存参数、优化器状态、step 和随机 key，才能真正恢复训练轨迹。

10. 下一步

到这里，JAX 101 关于表达计算的主线已经完整串起：数组与 jax.numpy、grad 与 vmap、用于组织结构的 PyTree、用于随机性的 key，以及用于状态的显式传值或 Ref。

让程序跑得更快需要继续学习 jax.jit、sharding 和 profiling。进入 MaxText 前，还要理解它如何把模型参数、优化器状态、指标、随机数和 checkpoint 组织成可变换的状态树。