第 6 章　JAX 中的状态管理

import sys
import warnings
from functools import partial
from typing import NamedTuple
from dataclasses import dataclass, field

import jax
import jax.numpy as jnp
from jax import lax

print("Python", sys.version.split()[0])
print("JAX   ", jax.__version__)
print("设备  ", jax.devices())

# 纯函数：输出只取决于输入，也不影响外界
def pure_add(a, b):
    return a + b

# 不纯：打印是一种 I/O 副作用
def impure_add_and_print(a, b):
    result = a + b
    print(f"结果是 {result}")   # 副作用
    return result

# 不纯：修改了全局变量
call_count = 0
def impure_increment_counter(x):
    global call_count
    call_count += 1             # 副作用：修改全局状态
    return x + 1

# 不纯：原地修改了输入参数
def impure_append_to_list(data_list, value):
    data_list.append(value)     # 副作用：修改了调用方的列表
    return data_list

print(pure_add(2, 3))
impure_add_and_print(2, 3)
impure_increment_counter(1); impure_increment_counter(1)
print("call_count =", call_count)
lst = [1, 2]
impure_append_to_list(lst, 3)
print("调用方的列表被改掉了：", lst)

# 典型的面向对象有状态写法：状态藏在对象内部，count() 每次都修改 self.n
class Counter:
    def __init__(self):
        self.n = 0

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

    def reset(self):
        self.n = 0

my_counter = Counter()
print(my_counter.count())   # 1
print(my_counter.count())   # 2  —— 同样的调用，不同的结果：不是纯函数

counter = 0

@jax.jit
def double(x):
    global counter
    counter += 1          # 副作用：只在追踪时执行
    return x * 2

for i in range(3):
    double(jnp.float32(i))
print("调用 3 次后 counter =", counter, "（只在第一次追踪时 +1）")
double(jnp.ones(3))       # 新形状 → 重新追踪
print("换一个形状再调用后 counter =", counter)
assert counter == 2

# 读全局变量：值在追踪时被固化
scale = 2.0
@jax.jit
def scaled(x):
    return x * scale

print("scale=2 时：", scaled(1.0))
scale = 10.0
print("改成 scale=10 后，标量输入（命中缓存）：", scaled(1.0))
print("改成 scale=10 后，新形状输入（重新编译）：", scaled(jnp.ones(2)))
assert float(scaled(1.0)) == 2.0 and float(scaled(jnp.ones(2))[0]) == 10.0   # 同一函数，新旧值并存

# 正确的调试打印方式
@jax.jit
def show(x):
    print("Python print：只在追踪时出现，x =", x)
    jax.debug.print("jax.debug.print：每次运行都出现，x = {}", x)
    return x + 1

show(1.0)
_ = show(2.0)

# 标准 Python：原地修改字典
my_data = {'values': [10, 20]}

def update_data(data_dict, new_value):
    data_dict['values'].append(new_value)   # 原地修改列表
    data_dict['last_added'] = new_value     # 原地新增键
    # 不需要 return，调用方的对象已经被改了

update_data(my_data, 30)
print(my_data)   # {'values': [10, 20, 30], 'last_added': 30}

# ❌ 错误写法：jit 一个原地修改对象属性的函数
class StateHolder:
    def __init__(self, value):
        self.value = jnp.array(value)

def update_state_impurely(state_obj, increment):
    # 不纯：原地修改对象的属性
    state_obj.value = state_obj.value + increment
    return state_obj.value

# 不加 JIT：按普通 Python 执行，看起来没问题
my_state = StateHolder(10.0)
print("非 JIT 第 1 次:", update_state_impurely(my_state, 5.0))   # 15.0
print("非 JIT 第 2 次:", update_state_impurely(my_state, 3.0))   # 18.0

# 加上 JIT：一定会报 TypeError
my_state_jit = StateHolder(10.0)
jitted_update = jax.jit(update_state_impurely)
stateholder_error = None
try:
    print("JIT 第 1 次:", jitted_update(my_state_jit, 5.0), my_state_jit.value)
except Exception as e:
    stateholder_error = e
    print(f"实测结果 → {type(e).__name__}: {str(e)[:200]}")

assert isinstance(stateholder_error, TypeError), "StateHolder 在 jit 下应当确定报 TypeError"

# ✅ 修正：显式状态传递 —— 状态是普通数组，函数返回新状态，不修改任何对象
def update_state_pure(value, increment):
    return value + increment

jitted_pure = jax.jit(update_state_pure)
v = jnp.array(10.0)
v = jitted_pure(v, 5.0); print("JIT 第 1 次:", v)
v = jitted_pure(v, 3.0); print("JIT 第 2 次:", v)
assert float(v) == 18.0

# 演示：注册为 PyTree 后，原地修改"静默失效"
@jax.tree_util.register_pytree_node_class
class StateHolderPytree:
    def __init__(self, value):
        self.value = value
    def tree_flatten(self):
        return (self.value,), None          # (子节点, 辅助数据)
    @classmethod
    def tree_unflatten(cls, aux, children):
        return cls(*children)

s = StateHolderPytree(jnp.array(10.0))
f = jax.jit(update_state_impurely)
r1 = f(s, 5.0); print("返回值:", r1, "  外部 s.value:", s.value)
r2 = f(s, 3.0); print("返回值:", r2, "  外部 s.value:", s.value, " ← 没有累加！静默错误")
assert float(s.value) == 10.0 and float(r2) == 13.0     # 外部对象从未被修改

# JAX 数组不可变
x = jnp.zeros(3)
try:
    x[0] = 1.0
except TypeError as e:
    print("TypeError:", str(e)[:110], "...")
print("函数式更新 x.at[0].set(1.0) →", x.at[0].set(1.0), "；原数组仍是", x)

# vmap 下的副作用：只在追踪时执行 1 次，拿到的是代表整批的 BatchTracer，并不存在"多个实例竞态"
side_effect_log = []
def append_and_double(v):
    side_effect_log.append(type(v).__name__)   # 副作用：修改外部列表
    return v * 2
jax.vmap(append_and_double)(jnp.arange(5.0))
print("vmap 处理 5 个样本后，外部列表被追加了", len(side_effect_log), "次：", side_effect_log)
assert side_effect_log == ["BatchTracer"]

# 命令式：状态藏在对象里
class Counter:
    def __init__(self):
        self.count = 0
    def increment(self):
        self.count += 1
        return self.count

c = Counter()
print(c.increment())   # 1
print(c.increment())   # 2
print("最终 count:", c.count)   # 2

# 函数式：状态显式传入、显式返回
def init_counter_state():
    """初始化状态。"""
    return 0

def increment(current_state):
    """接受当前状态，返回新状态。"""
    print(f"Running increment function with state: {current_state}")  # 用于演示
    new_state = current_state + 1
    return new_state

state = init_counter_state()
print("Initial state:", state)
state = increment(state)
print("State after first increment:", state)
state = increment(state)
print("State after second increment:", state)

jitted_increment = jax.jit(increment)

state = init_counter_state()
# 第一次调用触发追踪和编译，会执行 Python 代码（包括 print）
state = jitted_increment(state)
print("State after JIT call 1:", state)
# 之后的调用使用编译好的版本，内部 print 不再执行
state = jitted_increment(state)
print("State after JIT call 2:", state)
state = jitted_increment(state)
print("State after JIT call 3:", state)

params = {
    'layer1': {'weights': jnp.ones((3, 2)), 'bias': jnp.zeros((2,))},
    'layer2': {'weights': jnp.ones((2, 1)), 'bias': jnp.zeros((1,))},
}
print(jax.tree.map(lambda x: x.shape, params))

def update_params(params, grads, learning_rate):
    """对 PyTree 中的每个叶子执行一次 SGD 更新。"""
    def sgd_update(param, grad):
        return param - learning_rate * grad
    return jax.tree.map(sgd_update, params, grads)

# 占位梯度（结构与 params 相同）
grads = {
    'layer1': {'weights': jnp.full((3, 2), 0.1), 'bias': jnp.full((2,), 0.01)},
    'layer2': {'weights': jnp.full((2, 1), 0.2), 'bias': jnp.full((1,), 0.02)},
}
learning_rate = 0.01

updated_params = update_params(params, grads, learning_rate)
print("layer1 weights:\n", updated_params['layer1']['weights'])
print("layer2 bias:", updated_params['layer2']['bias'])

# 对 update_params 直接 jit，结构和数值与未编译版本一致
jit_update_params = jax.jit(update_params)
updated_jit = jit_update_params(params, grads, learning_rate)
same = jax.tree.all(jax.tree.map(jnp.allclose, updated_params, updated_jit))
print("jit 结果与非 jit 结果一致：", bool(same))
assert same

# 自证迁移表：顶层旧别名全部不存在，tree_util 版本和 jax.tree 版本都存在
removed_top_level = ["tree_map", "tree_leaves", "tree_structure", "tree_flatten",
                     "tree_unflatten", "tree_transpose", "tree_reduce"]
for name in removed_top_level:
    assert not hasattr(jax, name), f"jax.{name} 应已移除"
assert not hasattr(jax.tree_util, "tree_multimap")

still_in_tree_util = ["tree_map", "tree_leaves", "tree_structure", "tree_flatten", "tree_unflatten",
                      "tree_all", "tree_reduce", "tree_transpose", "tree_map_with_path",
                      "tree_leaves_with_path", "tree_broadcast"]
for name in still_in_tree_util:
    assert hasattr(jax.tree_util, name), f"jax.tree_util.{name} 应仍可用"

new_api = ["map", "leaves", "structure", "flatten", "unflatten", "all", "reduce",
           "transpose", "map_with_path", "leaves_with_path", "broadcast"]
for name in new_api:
    assert hasattr(jax.tree, name), f"jax.tree.{name} 应存在"
print("迁移表验证通过；jax.tree 模块提供：", [n for n in dir(jax.tree) if not n.startswith("_")])

# ❌ 旧写法（过时）：顶层别名 jax.tree_map 在 0.12 中已移除
try:
    jax.tree_map(lambda x: x * 2, {'a': 1.0})
except AttributeError as e:
    print("实测 → AttributeError:", e)

# ✅ 推荐写法：jax.tree.* 逐个演示
import operator
tree = {'b': jnp.arange(3.), 'a': (1.0, None, [2.0, 3.0])}

# (1) map：单棵树、多棵树（多棵树结构必须一致）
print("map 单棵树:", jax.tree.map(lambda x: x * 2, {'w': 1.0, 'b': (2.0, 3.0)}))
summed = jax.tree.map(lambda a, b, c: a + b + c, {'w': 1.0, 'b': 2.0}, {'w': 10.0, 'b': 20.0}, {'w': 100.0, 'b': 200.0})
print("map 三棵树:", summed)
assert summed == {'w': 111.0, 'b': 222.0}
try:
    jax.tree.map(lambda a, b: a + b, {'w': 1.0}, {'w': 1.0, 'b': 2.0})
except ValueError as e:
    print("结构不一致 → ValueError:", str(e).splitlines()[0])

# (2) leaves：None 是"空子树"不是叶子；dict 按键排序展平（'a' 在 'b' 之前）
lv = jax.tree.leaves(tree)
print("leaves:", lv)
assert len(lv) == 4 and lv[0] == 1.0

# (3) structure
print("structure:", jax.tree.structure(tree))

# (4) flatten / unflatten：往返重建
leaves, treedef = jax.tree.flatten(tree)
rebuilt = jax.tree.unflatten(treedef, [x * 10 for x in leaves])
print("unflatten:", rebuilt)
assert jax.tree.structure(rebuilt) == treedef

# (5) map_with_path：拿到每个叶子的键路径，例如只对 weights 做权重衰减
decayed = jax.tree.map_with_path(
    lambda path, x: x * 0.9 if 'weights' in jax.tree_util.keystr(path) else x, params)
print("map_with_path 路径:", jax.tree.map_with_path(lambda p, x: jax.tree_util.keystr(p), params))
assert jnp.allclose(decayed['layer1']['weights'], 0.9) and jnp.allclose(decayed['layer1']['bias'], 0.0)

# (6) reduce：把所有叶子归约成一个值，例如统计参数量、求梯度全局范数
n_params = jax.tree.reduce(operator.add, jax.tree.map(lambda x: x.size, params))
global_norm = jnp.sqrt(jax.tree.reduce(operator.add, jax.tree.map(lambda g: jnp.sum(g ** 2), grads)))
print("reduce 参数总数:", n_params, " 梯度全局范数:", global_norm)
assert n_params == 11

# (7) all：所有叶子都为真才为真
print("all:", jax.tree.all({'x': True, 'y': [True, True]}), jax.tree.all({'x': True, 'y': [True, False]}))
assert jax.tree.all({'x': True, 'y': [True, True]}) and not jax.tree.all({'x': True, 'y': [True, False]})

# (8) transpose："list of dict" ↔ "dict of list"，常用于把逐步记录的指标整理成按名字分组
records = [{'loss': 1.0, 'acc': 0.5}, {'loss': 0.8, 'acc': 0.6}, {'loss': 0.5, 'acc': 0.7}]
outer = jax.tree.structure([0] * len(records))
inner = jax.tree.structure({'loss': 0, 'acc': 0})
by_key = jax.tree.transpose(outer, inner, records)
print("transpose:", by_key)
assert by_key == {'acc': [0.5, 0.6, 0.7], 'loss': [1.0, 0.8, 0.5]}

# 1) NamedTuple
class TrainState(NamedTuple):
    step: jax.Array
    params: dict

ts = TrainState(step=jnp.array(0), params={'w': jnp.ones(2)})
ts2 = ts._replace(step=ts.step + 1)        # 生成新对象，原对象不变
print("NamedTuple 结构：", jax.tree.structure(ts))
print("ts.step =", ts.step, " ts2.step =", ts2.step)

# 2) register_dataclass：显式指定数据字段和元数据字段
@partial(jax.tree_util.register_dataclass, data_fields=['w', 'b'], meta_fields=['name'])
@dataclass(frozen=True)
class Linear:
    w: jax.Array
    b: jax.Array
    name: str          # 静态元数据，不是叶子

lin = Linear(jnp.ones(2), jnp.zeros(2), 'fc1')
print("dataclass 叶子：", jax.tree.leaves(lin))
print("dataclass 结构：", jax.tree.structure(lin))

# 也可以直接装饰，由字段元数据推断哪些是静态的
@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class AdamConfig:
    lr: jax.Array
    b1: float = field(default=0.9, metadata=dict(static=True))

print("推断结果：", jax.tree.structure(AdamConfig(jnp.array(1e-3))))

# frozen=True：尝试原地修改会立刻报错
try:
    lin.w = jnp.zeros(2)
except Exception as e:
    print(f"{type(e).__name__}: {e}")

# 这些容器都能直接穿过 jit / grad
@jax.jit
def forward(layer, x):
    return x @ layer.w + layer.b.sum()

print("jit 前向：", forward(lin, jnp.array([1.0, 2.0])))
print("对 dataclass 求梯度：", jax.grad(lambda l: forward(l, jnp.array([1.0, 2.0])))(lin))

class Counter:
    def __init__(self):
        self.count = 0
    def increment(self):
        self.count += 1     # 原地修改：副作用
        return self.count

counter_obj = Counter()
print(counter_obj.increment(), counter_obj.increment())

initial_state = 0

def update_counter(current_state):
    """将输入状态递增 1。"""
    print("Tracing update_counter...")   # 追踪时才会出现（演示用）
    return current_state + 1

state1 = update_counter(initial_state)
state2 = update_counter(state1)
state3 = update_counter(state2)
print("最终状态:", state3)

counter_traces = []
def update_counter_counted(current_state):
    counter_traces.append(1)          # 只在追踪时执行，用来数追踪次数
    return update_counter(current_state)

jitted_update_counter = jax.jit(update_counter_counted)

s = initial_state
for i in range(3):
    s = jitted_update_counter(s)
    print(f"第 {i + 1} 次调用后: {s}")
assert len(counter_traces) == 1 and int(s) == 3

def scan_body(carry_state, _):
    next_state = update_counter(carry_state)
    return next_state, None          # 每步没有额外输出

initial_scan_state = 0
num_steps = 10
final_state, accumulated_outputs = jax.lax.scan(
    scan_body, initial_scan_state, xs=None, length=num_steps)
print("scan 后的最终状态:", final_state)
print("每步输出:", accumulated_outputs)   # 每步输出 None，所以整体也是 None

# ❌ 错误写法：steps 作为普通参数传入 jit，再传给 scan 的 length
@jax.jit
def run_scan_jitted(init_state, steps):
    final, _ = jax.lax.scan(scan_body, init_state, xs=None, length=steps)
    return final

scan_error = None
try:
    print(run_scan_jitted(initial_scan_state, num_steps))
except Exception as e:
    scan_error = e
    print(f"实测结果 → {type(e).__name__}:\n{str(e)[:250]}")

assert isinstance(scan_error, jax.errors.ConcretizationTypeError)

# ✅ 修正 1：用 static_argnames 把 steps 标为静态（每个不同的 steps 值会各编译一次）
@partial(jax.jit, static_argnames="steps")
def run_scan_jitted(init_state, steps):
    final, _ = jax.lax.scan(scan_body, init_state, xs=None, length=steps)
    return final

r_static = run_scan_jitted(initial_scan_state, steps=num_steps)
print("修正 1（static_argnames）：", r_static)
assert int(r_static) == 10

# ✅ 修正 2：循环次数需要在运行时变化、又不需要逐步输出时，改用 fori_loop
@jax.jit
def run_fori(init_state, steps):
    return lax.fori_loop(0, steps, lambda i, c: update_counter(c), init_state)

r_fori_10, r_fori_25 = run_fori(initial_scan_state, 10), run_fori(initial_scan_state, 25)
print("修正 2（fori_loop）：", r_fori_10, r_fori_25)
assert int(r_fori_10) == 10 and int(r_fori_25) == 25

def running_mean_body(carry, x):
    new_carry = {'sum': carry['sum'] + x, 'n': carry['n'] + 1}
    y = new_carry['sum'] / new_carry['n']        # 每步输出：当前平均值
    return new_carry, y

xs = jnp.array([2.0, 4.0, 6.0, 8.0, 10.0])
init = {'sum': jnp.array(0.0), 'n': jnp.array(0)}
final, means = lax.scan(running_mean_body, init, xs)
print("最终状态:", final)
print("每步平均:", means)
assert jnp.allclose(means, jnp.array([2.0, 3.0, 4.0, 5.0, 6.0]))

# carry 的 dtype 不一致会报错
def add_body(carry, x):
    return carry + x, None          # int 状态 + float 输入 → float

try:
    lax.scan(add_body, jnp.array(0, dtype=jnp.int32), xs)    # 显式 int32 初始状态
except TypeError as e:
    print("dtype 不一致 → TypeError:", next(l for l in str(e).splitlines() if "input carry" in l)[:150])

# 例外：弱类型（Python 标量、jnp.array(0) 这种未指定 dtype 的值）允许被提升
print("弱类型初始值 jnp.array(0)：", lax.scan(add_body, jnp.array(0), xs)[0], "（未报错）")

# 形状不一致同样报错
try:
    lax.scan(lambda c, x: (jnp.stack([c, c]), None), jnp.array(0.0), xs)
except TypeError as e:
    print("形状不一致 → TypeError:", next(l for l in str(e).splitlines() if "input carry" in l)[:150])

def objective_function(x):
    return (x - 3.0) ** 2

grad_fn = jax.grad(objective_function)

def sgd_update(params, gradients, learning_rate):
    """纯函数：根据梯度返回新参数。"""
    return params - learning_rate * gradients

current_params = jnp.array(0.0)
learning_rate = 0.1
num_steps = 20

print(f"初始参数: {current_params:.4f}")
for step in range(num_steps):
    gradients = grad_fn(current_params)
    next_params = sgd_update(current_params, gradients, learning_rate)
    current_params = next_params          # 显式地把新状态接住
    if (step + 1) % 5 == 0:
        print(f"第 {step + 1:2d} 步: 参数 = {current_params:.4f}, 梯度 = {gradients:.4f}")
print(f"最终优化参数: {current_params:.4f}")
sgd_final = float(current_params)       # 记下来，稍后与动量法对比

# ❌ 超参数不当：momentum=0.9 在这个问题上欠阻尼，20 步内不收敛
class OptimizerState(NamedTuple):
    params: jax.Array
    velocity: jax.Array

def momentum_update(state, gradients, learning_rate, momentum):
    """纯函数：接收旧的优化器状态，返回新的优化器状态。"""
    new_velocity = momentum * state.velocity + learning_rate * gradients
    new_params = state.params - new_velocity
    return OptimizerState(params=new_params, velocity=new_velocity)

initial_params = jnp.array(0.0)
initial_velocity = jnp.array(0.0)
current_state = OptimizerState(params=initial_params, velocity=initial_velocity)

learning_rate = 0.1
momentum_coeff = 0.9
num_steps = 20

print(f"初始状态: 参数 = {current_state.params:.4f}, 速度 = {current_state.velocity:.4f}")
for step in range(num_steps):
    gradients = grad_fn(current_state.params)
    current_state = momentum_update(current_state, gradients, learning_rate, momentum_coeff)
    if (step + 1) % 5 == 0:
        print(f"第 {step + 1:2d} 步: 参数 = {current_state.params:.4f}, "
              f"速度 = {current_state.velocity:.4f}, 梯度 = {gradients:.4f}")
print(f"最终优化参数: {current_state.params:.4f}")
momentum09_final = float(current_state.params)
print(f"实测：|动量 μ=0.9 − 3| = {abs(momentum09_final - 3):.4f}  >  |普通 SGD − 3| = {abs(sgd_final - 3):.4f}")
assert abs(momentum09_final - 3) > abs(sgd_final - 3)     # μ=0.9 的动量法反而不如普通 SGD

# ✅ 修正：同样的 momentum_update，只把动量系数调到 0.5（学习率不变）
def run_momentum(mu, lr=0.1, steps=20):
    st = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
    for _ in range(steps):
        st = momentum_update(st, grad_fn(st.params), lr, mu)
    return float(st.params)

for mu_ in [0.0, 0.3, 0.5, 0.9]:
    print(f"momentum={mu_}: 20 步后参数 = {run_momentum(mu_):.4f}")

momentum05_final = run_momentum(0.5)
assert abs(momentum05_final - 3) < 1e-2                    # 调参后收敛
assert abs(momentum05_final - 3) < abs(sgd_final - 3)      # 并且优于普通 SGD

# ❌ 错误写法：学习率和动量设为静态参数（配合学习率调度时每步都重新编译）
momentum_traces = []
def momentum_update_logged(state, gradients, learning_rate, momentum):
    momentum_traces.append(learning_rate)        # 只在追踪时执行，用来数追踪次数
    return momentum_update(state, gradients, learning_rate, momentum)

jitted_momentum_update = jax.jit(momentum_update_logged, static_argnums=(2, 3))

schedule = [0.1, 0.09, 0.08, 0.07, 0.06]   # 模拟学习率衰减
current_state = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
print("--- 正在使用 JIT 编译的更新运行（静态学习率）---")
for lr_t in schedule:
    current_state = jitted_momentum_update(current_state, grad_fn(current_state.params), lr_t, 0.5)
print(f"参数 = {current_state.params:.4f}；追踪次数 = {len(momentum_traces)}")
assert len(momentum_traces) == 5

# ✅ 修正：学习率和动量作为普通（动态）参数传入
momentum_traces.clear()
jitted_momentum_update = jax.jit(momentum_update_logged)

current_state = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
print("--- 正在使用 JIT 编译的更新运行（动态学习率）---")
for lr_t in schedule:
    current_state = jitted_momentum_update(current_state, grad_fn(current_state.params), lr_t, 0.5)
print(f"参数 = {current_state.params:.4f}；追踪次数 = {len(momentum_traces)}")
assert len(momentum_traces) == 1

# 恒定学习率、跑满 20 步：结果与未编译版本一致
current_state = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
for step in range(num_steps):
    current_state = jitted_momentum_update(current_state, grad_fn(current_state.params), 0.1, 0.5)
print(f"最终优化参数 (JIT 编译, μ=0.5): {current_state.params:.4f}")
assert jnp.allclose(current_state.params, momentum05_final)

# 验证两种动量公式在学习率恒定时等价
g_fn = jax.grad(lambda x: (x - 3.0) ** 2)
eta, mu = 0.1, 0.9
x1, v1 = 0.0, 0.0      # Sutskever 写法
x2, v2 = 0.0, 0.0      # PyTorch / Optax 写法
for _ in range(20):
    g1 = g_fn(x1); v1 = mu * v1 + eta * g1; x1 = x1 - v1
    g2 = g_fn(x2); v2 = mu * v2 + g2;       x2 = x2 - eta * v2
print(f"Sutskever 写法: {float(x1):.6f}   PyTorch/Optax 写法: {float(x2):.6f}")
assert jnp.allclose(x1, x2, atol=1e-5)

# ---------- 通用优化器：接口仿照 Optax ----------
class Optimizer(NamedTuple):
    init: callable
    update: callable

def apply_updates(params, updates):
    return jax.tree.map(lambda p, u: p + u, params, updates)

class MomentumState(NamedTuple):
    velocity: dict

def sgd_momentum(lr, momentum=0.9):
    def init(params):
        return MomentumState(velocity=jax.tree.map(jnp.zeros_like, params))
    def update(grads, state, params=None):
        v = jax.tree.map(lambda v, g: momentum * v + g, state.velocity, grads)
        updates = jax.tree.map(lambda v: -lr * v, v)
        return updates, MomentumState(velocity=v)
    return Optimizer(init, update)

class AdamState(NamedTuple):
    count: jax.Array
    mu: dict
    nu: dict

def adam(lr, b1=0.9, b2=0.999, eps=1e-8):
    def init(params):
        # 注意：mu 和 nu 必须是两份独立的数组，不能共用同一个 zeros 对象（见 donate_argnums 补充）
        return AdamState(count=jnp.zeros([], jnp.int32),
                         mu=jax.tree.map(jnp.zeros_like, params),
                         nu=jax.tree.map(jnp.zeros_like, params))
    def update(grads, state, params=None):
        t = state.count + 1
        mu = jax.tree.map(lambda m, g: b1 * m + (1 - b1) * g, state.mu, grads)
        nu = jax.tree.map(lambda v, g: b2 * v + (1 - b2) * g ** 2, state.nu, grads)
        bc1 = 1 - b1 ** t                       # 偏差修正
        bc2 = 1 - b2 ** t
        updates = jax.tree.map(lambda m, v: -lr * (m / bc1) / (jnp.sqrt(v / bc2) + eps), mu, nu)
        return updates, AdamState(count=t, mu=mu, nu=nu)
    return Optimizer(init, update)

# ---------- 模型与数据 ----------
def init_mlp(key, sizes=(1, 32, 32, 1)):
    params = []
    keys = jax.random.split(key, len(sizes) - 1)
    for k, (n_in, n_out) in zip(keys, zip(sizes[:-1], sizes[1:])):
        params.append({'w': jax.random.normal(k, (n_in, n_out)) * jnp.sqrt(2.0 / n_in),
                       'b': jnp.zeros(n_out)})
    return params

def mlp(params, x):
    for layer in params[:-1]:
        x = jnp.tanh(x @ layer['w'] + layer['b'])
    return x @ params[-1]['w'] + params[-1]['b']

def loss_fn(params, x, y):
    return jnp.mean((mlp(params, x) - y) ** 2)

key = jax.random.key(0)
key, k_data, k_init = jax.random.split(key, 3)
X = jax.random.uniform(k_data, (256, 1), minval=-3.0, maxval=3.0)
Y = jnp.sin(X)
params0 = init_mlp(k_init)
print("参数结构：", jax.tree.map(jnp.shape, params0))

import numpy as np

# 在一个二次函数上自证：手写的 Adam / 动量 SGD 与 NumPy 手算参考实现逐步一致
A = np.array([1.0, 4.0, 0.25], dtype=np.float32)        # 各维曲率不同
C = np.array([3.0, -1.0, 2.0], dtype=np.float32)        # 最优点

def quad_loss(p):                                        # f(x) = 1/2 Σ a_i (x_i - c_i)^2
    return 0.5 * jnp.sum(A * (p['x'] - C) ** 2)

def numpy_adam_reference(x, steps, lr=0.1, b1=0.9, b2=0.999, eps=1e-8):
    m = np.zeros_like(x); v = np.zeros_like(x); traj = []
    for t in range(1, steps + 1):
        g = A * (x - C)                                  # 手算梯度
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g ** 2
        m_hat = m / (1 - b1 ** t); v_hat = v / (1 - b2 ** t)
        x = x - lr * m_hat / (np.sqrt(v_hat) + eps)
        traj.append(x.copy())
    return traj

def numpy_momentum_reference(x, steps, lr=0.1, mu=0.9):
    v = np.zeros_like(x); traj = []
    for _ in range(steps):
        g = A * (x - C)
        v = mu * v + g
        x = x - lr * v
        traj.append(x.copy())
    return traj

def run_jax_optimizer(opt, steps):
    p = {'x': jnp.zeros(3)}
    state = opt.init(p); traj = []
    for _ in range(steps):
        g = jax.grad(quad_loss)(p)
        updates, state = opt.update(g, state, p)
        p = apply_updates(p, updates)
        traj.append(np.asarray(p['x']))
    return traj, state

x0 = np.zeros(3, dtype=np.float32)
adam_traj, adam_state = run_jax_optimizer(adam(0.1), 3)
for t, (ours, ref) in enumerate(zip(adam_traj, numpy_adam_reference(x0, 3)), start=1):
    print(f"Adam 第 {t} 步  JAX: {ours}  NumPy: {ref}")
    assert np.allclose(ours, ref, rtol=1e-5, atol=1e-6)   # 约 1e-6 的差异来自 float32 计算 1-β2^t 的舍入
assert int(adam_state.count) == 3
# 偏差修正的效果：第 1 步每个坐标都移动约 lr=0.1，与梯度大小无关
assert np.allclose(np.abs(adam_traj[0]), 0.1, atol=1e-6)

mom_traj, _ = run_jax_optimizer(sgd_momentum(0.1, momentum=0.9), 3)
for ours, ref in zip(mom_traj, numpy_momentum_reference(x0, 3)):
    assert np.allclose(ours, ref, rtol=1e-5, atol=1e-6)
print("Adam 前 3 步、动量 SGD 前 3 步均与 NumPy 参考实现 allclose ✓")

def make_train(optimizer, num_steps):
    """返回 jit 编译的训练函数。参数、数据都作为函数参数传入（不靠闭包捕获大数组）；
    (params, opt_state) 作为 scan 的 carry。"""
    @jax.jit
    def train(params, x, y):
        def train_step(carry, _):
            params, opt_state = carry
            loss, grads = jax.value_and_grad(loss_fn)(params, x, y)   # x, y 是 train 的参数（tracer）
            updates, opt_state = optimizer.update(grads, opt_state, params)
            params = apply_updates(params, updates)
            return (params, opt_state), loss          # carry=状态；y=每步损失

        opt_state = optimizer.init(params)
        (params, opt_state), losses = lax.scan(train_step, (params, opt_state), None, length=num_steps)
        return params, opt_state, losses
    return train

results = {}
for name, opt in [("SGD", sgd_momentum(0.01, momentum=0.0)),
                  ("SGD+动量", sgd_momentum(0.01, momentum=0.9)),
                  ("Adam", adam(0.01))]:
    p, s, losses = make_train(opt, 1000)(params0, X, Y)
    results[name] = losses
    print(f"{name:8s} 初始损失 {losses[0]:.4f} → 最终损失 {losses[-1]:.5f}")

print("Adam 状态中的步数：", s.count, "；一阶矩结构与参数一致：",
      jax.tree.structure(s.mu) == jax.tree.structure(p))

import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['PingFang SC', 'Heiti SC', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

fig, ax = plt.subplots(figsize=(6, 3.5))
for name, losses in results.items():
    ax.plot(losses, label=name)
ax.set_yscale('log'); ax.set_xlabel('步数'); ax.set_ylabel('MSE 损失（对数刻度）')
ax.set_title('同一模型、三种优化器状态'); ax.legend()
plt.tight_layout(); plt.show()

opt = adam(1e-2)

@partial(jax.jit, donate_argnums=(0, 1))     # 捐赠 params 和 opt_state；数据 x, y 不捐赠
def train_step_donated(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = opt.update(grads, opt_state, params)
    return apply_updates(params, updates), opt_state, loss

params = jax.tree.map(jnp.copy, params0)
opt_state = opt.init(params)
old_leaf = params[0]['w']

params, opt_state, loss = train_step_donated(params, opt_state, X, Y)   # 正确：覆盖旧引用
print("旧参数缓冲区已被删除：", old_leaf.is_deleted())
assert old_leaf.is_deleted()
deleted_error = None
try:
    print(old_leaf + 1)
except RuntimeError as e:
    deleted_error = e
    print("再访问旧数组 → RuntimeError:", e)
assert deleted_error is not None
assert not X.is_deleted()                    # 没有捐赠的数据仍然可用

for _ in range(200):
    params, opt_state, loss = train_step_donated(params, opt_state, X, Y)
print(f"200 步后损失：{loss:.5f}")

# 无法利用的捐赠：输出里没有同形状的数组 → 警告而非报错
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    jax.jit(lambda a: a.sum(), donate_argnums=0)(jnp.ones(10))
msgs = [str(m.message) for m in w]
print("警告：", [m[:60] for m in msgs])
assert any("donated buffers were not usable" in m for m in msgs)

# ❌ 错误写法（同一 buffer 捐两次）：Adam 的 mu 和 nu 共用同一个 zeros 对象
def adam_init_shared(params):
    zeros = jax.tree.map(jnp.zeros_like, params)
    return AdamState(count=jnp.zeros([], jnp.int32), mu=zeros, nu=zeros)   # mu 与 nu 是同一批数组

params = jax.tree.map(jnp.copy, params0)
bad_opt_state = adam_init_shared(params)
double_donate_error = None
try:
    train_step_donated(params, bad_opt_state, X, Y)
except Exception as e:
    double_donate_error = e
    print(f"实测结果 → {type(e).__name__}: {str(e)[:120]}")
assert double_donate_error is not None and "donate the same buffer twice" in str(double_donate_error)
# （XLA 还会在 stderr 打印一行同样内容的 C++ 日志，可以忽略）

# ✅ 修正：mu 和 nu 各自创建（即前面 adam(...).init 的写法）
params = jax.tree.map(jnp.copy, params0)
good_opt_state = opt.init(params)
assert good_opt_state.mu[0]['w'] is not good_opt_state.nu[0]['w']
params, good_opt_state, loss = train_step_donated(params, good_opt_state, X, Y)
print(f"独立缓冲区 + 捐赠：正常运行，损失 = {loss:.5f}")

def stateful_counter(count_state, increment):
    """接收计数器状态和增量，返回新状态。"""
    new_count = count_state['count'] + increment
    return {'count': new_count}

initial_state = {'count': jnp.array(0)}

# 不使用 jit
state = stateful_counter(initial_state, 1)
state = stateful_counter(state, 2)
print("手动调用结果:", state)

# 使用 jit
jit_counter = jax.jit(stateful_counter)
jit_state = jit_counter(initial_state, 1)
jit_state = jit_counter(jit_state, 2)
print("JIT 调用结果:", jit_state)

# ❌ 旧写法（过时）：顶层别名 jax.tree_map 在 JAX 0.12 中已移除
tree_map_error = None
try:
    assert jax.tree_util.tree_all(jax.tree_map(lambda x, y: jnp.array_equal(x, y), state, jit_state))
except AttributeError as e:
    tree_map_error = e
    print("实测结果 → AttributeError:", e)
assert tree_map_error is not None

# ✅ 推荐写法
assert jax.tree.all(jax.tree.map(lambda x, y: jnp.array_equal(x, y), state, jit_state))
print("jax.tree.map 断言通过：两种结果一致")

traces = []
@jax.jit
def bump(state):
    traces.append(jax.tree.structure(state))     # 追踪时记录
    return jax.tree.map(lambda x: x + 1, state)

bump({'count': jnp.array(0)})
bump({'count': jnp.array(5)})                     # 结构相同：命中缓存
bump({'count': jnp.array(0), 'steps': jnp.array(0)})   # 新增键：重新追踪
print("追踪次数:", len(traces))
for t in traces: print("  ", t)
assert len(traces) == 2

def predict_and_update(params, x):
    """前向预测；参数状态原样返回（此处不修改）。"""
    pred = params['w'] * x + params['b']
    return pred, params

def loss_fn_linear(params, x, y_target):
    pred, _ = predict_and_update(params, x)
    return jnp.mean((pred - y_target) ** 2)

# 数据：理想参数为 w=2.5, b=1.0
x_data = jnp.array([1.0, 2.0, 3.0])
y_data = jnp.array([3.5, 5.5, 7.5])
params = {'w': jnp.array(2.0), 'b': jnp.array(1.0)}

grad_fn_linear = jax.grad(loss_fn_linear, argnums=0)
grads = grad_fn_linear(params, x_data, y_data)
print("梯度（与 params 同结构）:", grads)

loss_val, grads = jax.value_and_grad(loss_fn_linear)(params, x_data, y_data)
print("损失:", loss_val, " 梯度:", grads)