第 1 章　高级变换与控制流

# 公共导入与工具函数
import time, io, contextlib, functools
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax

jax.config.update("jax_pprint_use_color", False)   # 打印 jaxpr 时不带 ANSI 颜色码，便于阅读和字符串匹配
print("jax", jax.__version__, "| 设备:", jax.devices())

N_ASSERT = 0  # 统计本 notebook 中通过的 assert 数量
def check(cond, msg=""):
    """包装 assert：失败就抛错，成功就计数。"""
    global N_ASSERT
    assert cond, msg
    N_ASSERT += 1

def bench(f, *args, n=20):
    """计时工具：先预热（包括编译），再取 n 次的中位数，单位为毫秒。"""
    jax.block_until_ready(f(*args))
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args))
        ts.append(time.perf_counter() - t0)
    return 1e3 * float(np.median(ts))

def slow_f(x):
    return jnp.sin(x) * x + jnp.log(x + 1)

fast_f = jax.jit(slow_f)

key = jax.random.key(0)                      # 0.12 推荐：带类型的 key
x = jax.random.uniform(key, (1000, 1000))    # 取 [0,1) 区间，保证 log(x+1) 有定义

y_eager = slow_f(x).block_until_ready()
y_jit = fast_f(x).block_until_ready()        # 第一次调用：追踪 + 编译
print("结果一致:", bool(jnp.allclose(y_eager, y_jit)))
check(jnp.allclose(y_eager, y_jit))

print(f"逐操作执行（eager）: {bench(slow_f, x):.3f} ms | jit: {bench(fast_f, x):.3f} ms")

# ❌ 旧写法（过时）：jax.random.PRNGKey 生成的是“裸”uint32 数组形式的旧式 key
try:
    old_key = jax.random.PRNGKey(0)
    print("PRNGKey 仍可运行 → dtype =", old_key.dtype, ", shape =", old_key.shape)
except Exception as e:
    print("报错:", type(e).__name__, e)

# ✅ 推荐写法：jax.random.key 生成带类型的 key（dtype 为 key<fry>，shape 为 ()）
new_key = jax.random.key(0)
print("dtype =", new_key.dtype, ", shape =", new_key.shape)
# 两者底层的位数据相同，只是新式 key 不会被误当成普通整数数组参与运算
check(np.array_equal(jax.random.key_data(new_key), jax.random.PRNGKey(0)))

def predict(params, x):
    w, b = params
    return jnp.dot(w, x) + b

def loss_fn(params, x, y_true):
    y_pred = predict(params, x)
    return jnp.mean((y_pred - y_true) ** 2)

w_init = jnp.array([1.5, -0.5]); b_init = 0.3
params_init = (w_init, b_init)
x_data = jnp.array([0.2, 0.8]); y_target = 2.5

grad_loss_fn = jax.grad(loss_fn, argnums=0)
gw, gb = grad_loss_fn(params_init, x_data, y_target)
print("关于 w 的梯度:", gw)
print("关于 b 的梯度:", gb)

# 用 NumPy float64 按解析公式对拍
w64, x64 = np.array([1.5, -0.5]), np.array([0.2, 0.8])
r = w64 @ x64 + 0.3 - 2.5            # 残差 = 0.2 - 2.5 = -2.3
check(np.allclose(gw, 2 * r * x64) and np.isclose(gb, 2 * r))
check(np.allclose(gw, [-0.92, -3.68]) and np.isclose(gb, -4.6))

def simple_affine(w, b, x):
    return jnp.dot(w, x) + b

w = jnp.array([[1., 2.], [3., 4.], [5., 6.]])   # (3, 2)
b = jnp.array([0.1, 0.2, 0.3])                  # (3,)
x_batch = jnp.arange(8.).reshape(4, 2)          # (4, 2)

batched_affine = jax.vmap(simple_affine, in_axes=(None, None, 0))
out = batched_affine(w, b, x_batch)
print("输出形状:", out.shape)
check(out.shape == (4, 3))
check(np.allclose(out, np.asarray(x_batch) @ np.asarray(w).T + np.asarray(b)))

key = jax.random.key(0)
key, kx, ky = jax.random.split(key, 3)          # 一次切出需要的所有子 key，不复用
x_batch_data = jax.random.normal(kx, (16, 2))
y_batch_target = jax.random.normal(ky, (16,))

batched_grad_loss_fn = jax.jit(jax.vmap(grad_loss_fn, in_axes=(None, 0, 0)))
gw_b, gb_b = batched_grad_loss_fn(params_init, x_batch_data, y_batch_target)
print("w 的逐样本梯度形状:", gw_b.shape, "| b 的逐样本梯度形状:", gb_b.shape)
check(gw_b.shape == (16, 2) and gb_b.shape == (16,))

def mean_loss(params, xs, ys):
    return jnp.mean(jax.vmap(loss_fn, in_axes=(None, 0, 0))(params, xs, ys))

gw_mean, gb_mean = jax.jit(jax.grad(mean_loss))(params_init, x_batch_data, y_batch_target)
print("平均损失的梯度:", gw_mean, gb_mean)
check(np.allclose(gw_mean, gw_b.mean(0), atol=1e-6) and np.allclose(gb_mean, gb_b.mean(), atol=1e-6))

def py_loop_sum(x, n):
    s = 0.0
    for _ in range(n):       # n 必须是 Python 整数
        s = s + x
    return s

# 情况 1：n 是 traced 值 → 报错
try:
    jax.jit(py_loop_sum)(1.0, 3)
    raised = None
except Exception as e:
    raised = type(e).__name__
print("n 为 traced 值:", raised)
check(raised == "TracerIntegerConversionError")

# 情况 2：n 是静态值 → 展开，jaxpr 中的 add 数量随 n 线性增长
n_eqns = {n: len(jax.make_jaxpr(py_loop_sum, static_argnums=1)(1.0, n).eqns) for n in (3, 30)}
print("静态 n 时 jaxpr 的方程数:", n_eqns)
check(n_eqns[30] > n_eqns[3] * 5)

def cumulative_sum_step(carry, x):
    new_carry = carry + x
    return new_carry, new_carry          # 输出“截至当前”的累加和

xs = jnp.array([1, 2, 3, 4, 5])
final, ys = lax.scan(cumulative_sum_step, 0, xs)
print("最终 carry:", final, "| 序列:", ys)
check(int(final) == 15 and ys.tolist() == [1, 3, 6, 10, 15])

# ❌ 旧写法（错误）：W_hh 与 W_xh 共用同一个 w_key；x_key = random.split(key) 得到的是 key 数组
from jax import random
key = random.PRNGKey(0)
key, w_key, b_key = random.split(key, 3)
W_hh = random.normal(w_key, (5, 5)) * 0.1
W_xh = random.normal(w_key, (5, 3)) * 0.1      # 与 W_hh 用的是同一个 key！
# 运行可见：W_xh 的 15 个元素恰好就是 W_hh 展平后的前 15 个元素，两者根本不独立
same = bool(np.allclose(np.ravel(W_xh), np.ravel(W_hh)[:15]))
print("W_xh == W_hh 展平后的前 15 个元素:", same)
check(same)
try:
    x_key = random.split(key)                    # 形状 (2, 2) 的 key 数组，不是单个 key
    input_sequence = random.normal(x_key, (4, 3))
    print("未报错")
except Exception as e:
    print("报错:", type(e).__name__, "-", str(e)[:90])

# ✅ 推荐写法：jax.random.key + 每个用途一个独立子 key；参数作为函数参数传入（不靠闭包）
input_features, hidden_features, sequence_length = 3, 5, 4
k_hh, k_xh, k_b, k_x = jax.random.split(jax.random.key(0), 4)
rnn_params = {
    "W_hh": jax.random.normal(k_hh, (hidden_features, hidden_features)) * 0.1,
    "W_xh": jax.random.normal(k_xh, (hidden_features, input_features)) * 0.1,
    "b_h":  jax.random.normal(k_b, (hidden_features,)) * 0.1,
}
input_sequence = jax.random.normal(k_x, (sequence_length, input_features))
h_initial = jnp.zeros((hidden_features,))

def rnn_step(params, h_prev, x_t):
    h_new = jnp.tanh(params["W_hh"] @ h_prev + params["W_xh"] @ x_t + params["b_h"])
    return h_new, h_new                       # (新 carry, 本步输出)

@jax.jit
def run_rnn_scan(params, h0, xs):
    # 在 jit 内部，params 已经是 tracer，lambda 捕获它不会产生嵌入常量
    return lax.scan(lambda h, x: rnn_step(params, h, x), h0, xs)

h_final, h_seq = run_rnn_scan(rnn_params, h_initial, input_sequence)
print("输入:", input_sequence.shape, "初始:", h_initial.shape, "最终:", h_final.shape, "序列:", h_seq.shape)
check(h_final.shape == (5,) and h_seq.shape == (4, 5))

# 与 NumPy float64 的 Python 循环对拍
P = {k: np.asarray(v, np.float64) for k, v in rnn_params.items()}
h = np.zeros(5); ref = []
for x_t in np.asarray(input_sequence, np.float64):
    h = np.tanh(P["W_hh"] @ h + P["W_xh"] @ x_t + P["b_h"]); ref.append(h)
check(np.allclose(h_seq, np.stack(ref), atol=1e-6) and np.allclose(h_final, h, atol=1e-6))
print("与 Python 循环结果一致 ✓")

def generate_powers_of_2(carry, _):          # 第二个参数会收到 None
    return carry * 2, carry                  # 新 carry 翻倍，输出当前值

final_val, powers = lax.scan(generate_powers_of_2, 1, xs=None, length=10)
print("最终值:", final_val, "| 序列:", powers)
check(int(final_val) == 1024 and powers.tolist() == [2 ** i for i in range(10)])

# ❌ 错误写法：以为 xs 为 None 时，循环体可以只写成 f(carry)
try:
    lax.scan(lambda c: (c * 2, c), 1, xs=None, length=3)
    print("未报错")
except Exception as e:
    print("报错:", type(e).__name__, "-", e)

# ✅ 推荐写法：循环体始终接收两个参数，不用的那个写成 _
final_val, _ = lax.scan(lambda c, _: (c * 2, None), 1, xs=None, length=3)   # y 为 None 时 ys 也是 None
check(int(final_val) == 8)
try:
    lax.scan(lambda c: (c * 2, c), 1, xs=None, length=3); ok = False
except TypeError:
    ok = True
check(ok)
print("f(carry) 形式会抛 TypeError；f(carry, _) 正常 ✓")

f = lambda c, x: (c + x, c + x)
xs = jnp.arange(1, 6)
_, ys_rev = lax.scan(f, 0, xs, reverse=True)
_, ys_flip = lax.scan(f, 0, xs[::-1])
print("reverse=True:", ys_rev, "（后缀和，ys[t] 对应 xs[t]）")
check(ys_rev.tolist() == [15, 14, 12, 9, 5])
check(np.array_equal(ys_rev, ys_flip[::-1]))

def make_scan(unroll):
    @jax.jit
    def run(x0, xs):
        return lax.scan(lambda c, x: (jnp.sin(c) * 0.9 + x, None), x0, xs, unroll=unroll)[0]
    return run

xs_long = jnp.linspace(0., 1., 10_000)
results = {}
for u in (1, 4, 16, True):
    fn = make_scan(u)
    results[u] = fn(0.0, xs_long)
    print(f"unroll={str(u):>4}: {bench(fn, 0.0, xs_long):7.3f} ms")
check(all(np.allclose(results[1], v, atol=1e-5) for v in results.values()))   # 语义不变
# jaxpr 中可以看到 unroll 参数被记录下来
print([l.strip() for l in str(jax.make_jaxpr(make_scan(4))(0.0, xs_long)).splitlines() if "unroll" in l])

def py_loop(x):
    for _ in range(500):
        x = jnp.sin(x) + 0.1
    return x

def scan_loop(x):
    return lax.scan(lambda c, _: (jnp.sin(c) + 0.1, None), x, None, length=500)[0]

for fn in (py_loop, scan_loop):
    t0 = time.perf_counter(); jax.jit(fn).lower(1.0).compile()
    print(f"{fn.__name__:>9}: 追踪+编译 {1e3 * (time.perf_counter() - t0):7.1f} ms, "
          f"jaxpr 方程数 = {len(jax.make_jaxpr(fn)(1.0).eqns)}")
check(np.allclose(jax.jit(py_loop)(1.0), jax.jit(scan_loop)(1.0), atol=1e-5))

try:
    lax.scan(lambda c, x: (c + x, None), 0, jnp.arange(5), length=4); msg = None
except ValueError as e:
    msg = str(e)
print("length 不一致:", msg)
check(msg is not None and "disagrees" in msg)

# pytree xs / ys：字典输入，元组输出
xs_tree = {"a": jnp.arange(4.), "b": jnp.ones((4, 2))}
def body(c, x):
    c = c + x["a"] * x["b"]             # carry 形状 (2,)
    return c, (c.sum(), x["a"] ** 2)
c_final, (sums, sq) = lax.scan(body, jnp.zeros(2), xs_tree)
print("final:", c_final, "| sums:", sums, "| sq:", sq)
check(np.allclose(c_final, [6., 6.]) and np.allclose(sq, [0, 1, 4, 9]))

def scan_err(init, body, xs):
    try:
        lax.scan(body, init, xs); return "（未报错）"
    except TypeError as e:
        return str(e).split("\n\n")[1]       # 取出中间那句具体说明

xs_f = jnp.arange(5.)
m_dtype = scan_err(jnp.int32(0), lambda c, x: (c + x, c), xs_f)
m_shape = scan_err(jnp.float32(0), lambda c, x: (jnp.stack([c, c]), c), xs_f)
m_tree  = scan_err(jnp.float32(0), lambda c, x: ((c, c), c), xs_f)
for name, m in [("dtype", m_dtype), ("形状", m_shape), ("结构", m_tree)]:
    print(f"[{name}] {m}")
check("dtypes do not match" in m_dtype and "shapes do not match" in m_shape and "Python types differ" in m_tree)

# 弱类型 init 不报错，carry 被提升为 float32
c_weak, _ = lax.scan(lambda c, x: (c + x, None), 0, xs_f)
print("init=0（弱类型）→ final carry dtype:", c_weak.dtype)
check(c_weak.dtype == jnp.float32)

x = jnp.arange(1., 9.)
check(np.allclose(lax.associative_scan(jnp.add, x), jnp.cumsum(x)))

# 线性递推 h_t = a_t * h_{t-1} + b_t，二元组合算子 (a1,b1)∘(a2,b2) = (a1*a2, a2*b1 + b2)
a = jnp.full(8, 0.9); bb = jnp.arange(8.)
comb = lambda l, r: (l[0] * r[0], r[0] * l[1] + r[1])
_, h_par = lax.associative_scan(comb, (a, bb))
_, h_seq = lax.scan(lambda h, ab: (ab[0] * h + ab[1],) * 2, 0.0, (a, bb))
print("并行前缀:", h_par)
check(np.allclose(h_par, h_seq, atol=1e-5))

@jax.jit
def python_conditional_process(x):
    if jnp.sum(x) > 0:          # 需要把 tracer 转成 bool
        return x * x
    return x * x * x

for data in (jnp.array([1., 2., 3.]), jnp.array([-1., -2., -3.])):
    try:
        python_conditional_process(data); err = None
    except Exception as e:
        err = e
    print(type(err).__name__, "| 是 ConcretizationTypeError 的子类:",
          isinstance(err, jax.errors.ConcretizationTypeError))
check(isinstance(err, jax.errors.TracerBoolConversionError))

def lax_conditional_process(x):
    def true_branch(operand):
        print("  [追踪] true_branch")
        return operand * operand
    def false_branch(operand):
        print("  [追踪] false_branch")
        return operand * operand * operand
    pred = jnp.sum(x) > 0
    return lax.cond(pred, true_branch, false_branch, x)

jit_cond = jax.jit(lax_conditional_process)
print("第一次调用（正数）:")
r_pos = jit_cond(jnp.array([1., 2., 3.]))
print("  结果:", r_pos)
print("第二次调用（负数，形状相同 → 复用编译结果，不再打印追踪信息）:")
r_neg = jit_cond(jnp.array([-1., -2., -3.]))
print("  结果:", r_neg)
check(np.allclose(r_pos, [1, 4, 9]) and np.allclose(r_neg, [-1, -8, -27]))

print("多操作数:", lax.cond(True, lambda a, b: a + b, lambda a, b: a - b, 1., 2.))

def cond_err(*args):
    try:
        lax.cond(*args); return None
    except TypeError as e:
        return str(e).splitlines()[0]
m1 = cond_err(True, lambda x: x, lambda x: x.astype(jnp.int32), 1.)
m2 = cond_err(jnp.array([True, False]), lambda x: x, lambda x: -x, 1.)
print("分支 dtype 不一致:", m1)
print("pred 非标量:", m2)
check(m1.startswith("cond branches must have equal output types") and m2.startswith("Pred must be a scalar"))

branches = [lambda x: x + 1, lambda x: x * 2, lambda x: -x]
out = [float(lax.switch(i, branches, 3.0)) for i in (0, 1, 2, 5, -3)]
print("switch index 0,1,2,5,-3 →", out)
check(out == [4.0, 6.0, -3.0, -3.0, 4.0])          # 5 被截断成 2，-3 被截断成 0

z, o, t = jnp.zeros(4), jnp.ones(4), 2 * jnp.ones(4)
s_int = lax.select_n(jnp.array([0, 1, 2, 1]), z, o, t)
s_bool = lax.select_n(jnp.array([True, False, True, False]), z, o)
print("select_n(int) :", s_int)
print("select_n(bool):", s_bool, "← True 选第 1 个 case（这里是 ones）")
check(np.allclose(s_int, [0, 1, 2, 1]) and np.allclose(s_bool, [1, 0, 1, 0]))
# 与 jnp.where 的顺序对照：where(c, x, y) == select_n(c, y, x)
c = jnp.array([True, False, True, False])
check(np.allclose(jnp.where(c, o, z), lax.select_n(c, z, o)))

def g(x):
    return lax.cond(x > 0, lambda v: v * 2, lambda v: -v, x)

jp_scalar = str(jax.make_jaxpr(g)(1.0))
jp_batched_pred = str(jax.make_jaxpr(jax.vmap(g))(jnp.arange(3.)))
print("—— vmap 后（pred 被批处理）的 jaxpr ——\n", jp_batched_pred)
check("cond[" in jp_scalar)
check("cond[" not in jp_batched_pred and "select_n" in jp_batched_pred)

def g2(p, x):
    return lax.cond(p > 0, lambda v: v * 2, lambda v: -v, x)
jp_unbatched_pred = str(jax.make_jaxpr(jax.vmap(g2, in_axes=(None, 0)))(1.0, jnp.arange(3.)))
check("cond[" in jp_unbatched_pred)
print("pred 不带批维时，cond 原语保留 ✓")

def find_power_of_two(n):
    def cond_fun(current_power):
        return current_power < n
    def body_fun(current_power):
        return current_power * 2
    init_val = 1
    return lax.while_loop(cond_fun, body_fun, init_val)

target_number = 100
r1 = find_power_of_two(target_number)
r2 = jax.jit(find_power_of_two)(target_number)
print("最小的 2 的幂 >= 100:", r1, "| JIT 结果:", r2)
check(int(r1) == 128 and int(r2) == 128)
# 循环轨迹：cond(1)…cond(64) 为真，body(64)=128，cond(128) 为假 → 返回 128

def power_loop(x):
    # 计算 x**5：计数器 i 从 0 到 5
    return lax.while_loop(lambda c: c[0] < 5, lambda c: (c[0] + 1, c[1] * x), (0, 1.0))[1]

val, tangent = jax.jvp(power_loop, (2.0,), (1.0,))
print("x^5 at 2 =", val, "| 前向模式导数 5x^4 =", tangent)
check(np.isclose(val, 32.) and np.isclose(tangent, 80.))

try:
    jax.grad(power_loop)(2.0); msg = None
except ValueError as e:
    msg = str(e)
print("grad(while_loop):", msg)
check(msg is not None and "Reverse-mode differentiation does not work" in msg)

# jacfwd（前向模式构造的雅可比）可以用
check(np.isclose(jax.jacfwd(power_loop)(2.0), 80.))

def fl_static(x):
    return lax.fori_loop(0, 5, lambda i, c: c * x, 1.0)
def fl_dynamic(x, n):
    return lax.fori_loop(0, n, lambda i, c: c * x, 1.0)

jp_s = str(jax.make_jaxpr(fl_static)(2.0))
jp_d = str(jax.make_jaxpr(fl_dynamic)(2.0, 5))
print("静态上下界 → 原语:", "scan" if "scan[" in jp_s else "while")
print("动态上下界 → 原语:", "scan" if "scan[" in jp_d else "while")
check("scan[" in jp_s and "while[" in jp_d)

check(np.isclose(jax.grad(fl_static)(2.0), 80.))                 # 静态 → 可以反向求导
try:
    jax.jit(jax.grad(fl_dynamic))(2.0, 5); ok = False             # jit 下 n 是 traced 值 → 动态
except ValueError:
    ok = True
check(ok)
print("fori_loop：静态上下界可 grad，traced 上下界不可 ✓")

def collatz_steps(n):
    def body(c):
        v, k = c
        return jnp.where(v % 2 == 0, v // 2, 3 * v + 1), k + 1
    return lax.while_loop(lambda c: c[0] > 1, body, (n, 0))[1]

ns = jnp.array([1, 2, 3, 6, 7, 27])
steps = jax.vmap(collatz_steps)(ns)
print("各元素的 Collatz 步数:", steps, "→ 整个批次要循环", int(steps.max()), "次")

def collatz_py(n):
    k = 0
    while n > 1:
        n = n // 2 if n % 2 == 0 else 3 * n + 1; k += 1
    return k
check(steps.tolist() == [collatz_py(int(n)) for n in ns])

hlo = jax.jit(jax.vmap(collatz_steps)).lower(ns).as_text()
check("stablehlo.or" in hlo and "stablehlo.select" in hlo)
print("HLO 中可以看到 reduce-or（“任一元素未完成”）和 select（冻结已完成元素）✓")

def conditional_computation(x, threshold=5.0):
    total = jnp.sum(x)
    return lax.cond(total >= threshold, lambda o: o * 2.0, lambda o: o / 2.0, x)

jit_cc = jax.jit(conditional_computation)
r1 = jit_cc(jnp.array([1., 2., 3.]))     # 和为 6 ≥ 5 → ×2
r2 = jit_cc(jnp.array([1., 1., 1.]))     # 和为 3 < 5 → ÷2
print("Result 1:", r1, "| Result 2:", r2)
check(np.allclose(r1, [2, 4, 6]) and np.allclose(r2, [0.5, 0.5, 0.5]))

def simple_rnn_step(carry, x_t):
    W, U, b = 0.5, 0.8, 0.1                 # Python 标量常量：作为字面量嵌入，没有问题
    new_hidden = jnp.tanh(W * x_t + U * carry + b)
    return new_hidden, new_hidden

def run_rnn(initial_hidden, inputs):
    _, outputs = lax.scan(simple_rnn_step, initial_hidden, inputs)
    return jnp.sum(outputs)

grad_run_rnn = jax.grad(run_rnn, argnums=0)
input_sequence_s = jnp.array([1.0, -0.5, 2.0])
g_h0 = grad_run_rnn(0.0, input_sequence_s)
print("dL/dh0 =", g_h0)

# NumPy float64 对拍：手动推导公式 + 中心差分
xs64 = np.array([1.0, -0.5, 2.0])
def rnn_np(h0):
    h, hs = h0, []
    for x_t in xs64:
        h = np.tanh(0.5 * x_t + 0.8 * h + 0.1); hs.append(h)
    return np.array(hs)
hs = rnn_np(0.0)
analytic = sum(np.prod(0.8 * (1 - hs[:t + 1] ** 2)) for t in range(3))
fd = (rnn_np(1e-6).sum() - rnn_np(-1e-6).sum()) / 2e-6
print(f"手推公式: {analytic:.7f} | 中心差分: {fd:.7f}")
check(np.isclose(g_h0, analytic, atol=1e-5) and np.isclose(analytic, fd, atol=1e-6))
check(abs(float(g_h0) - 0.313) > 0.5)        # 和常见的错误数值 0.313 明显不同

def conditional_loss(params, x):
    pred = jnp.sum(x) > 0
    def loss1(p): return jnp.sum(p * x)
    def loss2(p): return jnp.sum(p / (x + 1e-5))
    return lax.cond(pred, loss1, loss2, params)

params = jnp.array([0.5, -0.5])
g_pos = jax.grad(conditional_loss)(params, jnp.array([1.0, 1.0]))
g_neg = jax.grad(conditional_loss)(params, jnp.array([-1.0, -1.0]))
print("正例梯度:", g_pos, "| 负例梯度:", g_neg, "| 理论值 1/(-1+1e-5) =", 1 / (-1 + 1e-5))
check(np.allclose(g_pos, [1., 1.]))
check(np.allclose(g_neg, 1 / (-1 + 1e-5), rtol=1e-6) and float(g_neg[0]) < -1.0)

batched_run_rnn = jax.vmap(run_rnn, in_axes=(0, 0))
batch_hidden_init = jnp.zeros(4)
batch_input_sequence = jnp.arange(12, dtype=jnp.float32).reshape((4, 3))
out_b = batched_run_rnn(batch_hidden_init, batch_input_sequence)
print("批量输出形状:", out_b.shape)
check(out_b.shape == (4,))
check(np.allclose(out_b, [run_rnn(h, s) for h, s in zip(batch_hidden_init, batch_input_sequence)]))

vmapped_conditional = jax.vmap(conditional_computation, in_axes=(0, None))
batch = jnp.array([[1., 2., 3.], [1., 1., 1.], [10., 1., 1.]])
out_v = vmapped_conditional(batch, 5.0)
print(out_v)
check(np.allclose(out_v, [[2, 4, 6], [0.5, 0.5, 0.5], [20, 2, 2]]))

# ❌ 旧写法（错误）：grad 直接套在 vmap 外面，但 vmap 的输出形状是 (2, 3)，不是标量
def threshold_from_params(params, x):
    threshold = jnp.mean(params)
    return conditional_computation(x, threshold)

batch_data = jnp.array([[1., 2., 3.], [1., 1., 1.]])
params_t = jnp.array([4.0, 6.0])                  # 均值为 5.0
grad_vmap_fn = jax.jit(jax.grad(jax.vmap(threshold_from_params, in_axes=(None, 0)), argnums=0))
try:
    print(grad_vmap_fn(params_t, batch_data))
except Exception as e:
    print("报错:", type(e).__name__, "-", str(e)[:110])

# ✅ 推荐写法：先把 vmap 的输出归约成标量，再求 grad
def total_output(params, xs):
    return jnp.sum(jax.vmap(threshold_from_params, in_axes=(None, 0))(params, xs))

g_thr = jax.jit(jax.grad(total_output))(params_t, batch_data)
print("对阈值参数的梯度:", g_thr)
check(np.allclose(g_thr, 0.0))            # 阈值只通过比较运算进入 pred → 梯度恒为 0

try:
    jax.grad(jax.vmap(threshold_from_params, in_axes=(None, 0)))(params_t, batch_data); ok = False
except TypeError as e:
    ok = "scalar-output" in str(e)
check(ok)

padded_sequences = jnp.array([[1, 2, 3, 0], [4, 5, 0, 0], [6, 7, 8, 9]], dtype=jnp.float32)
mask = padded_sequences != 0
print(mask)

# 例 1：掩码求和
masked_sum = jnp.sum(jnp.where(mask, padded_sequences, 0), axis=-1)
# 例 2：掩码平均
num_valid_elements = jnp.sum(mask, axis=-1)
masked_avg = masked_sum / jnp.maximum(num_valid_elements, 1)
print("求和:", masked_sum, "| 平均:", masked_avg)
check(np.allclose(masked_sum, [6, 9, 30]) and np.allclose(masked_avg, [2.0, 4.5, 7.5]))

lengths = jnp.array([3, 2, 4])
mask_from_len = jnp.arange(4)[None, :] < lengths[:, None]
check(np.array_equal(mask_from_len, mask))
print(mask_from_len)

# ❌ 旧写法（错误）：jnp.randn 不是 JAX API（JAX 的随机数必须显式传入 key）
try:
    scores = jnp.randn(3, 4)
except AttributeError as e:
    print("报错:", type(e).__name__, "-", e)

# ✅ 推荐写法：jax.random.normal + 显式 key
scores = jax.random.normal(jax.random.key(0), (3, 4))
masked_scores = jnp.where(mask, scores, -1e9)
attention_weights = jax.nn.softmax(masked_scores, axis=-1)
print(attention_weights)
check(np.allclose(np.asarray(attention_weights)[~np.asarray(mask)], 0.0))
check(np.allclose(attention_weights.sum(-1), 1.0, atol=1e-6))

# 0.12 中 softmax 自带 where= 参数，效果相同，且被屏蔽的位置输出恰为 0
aw2 = jax.nn.softmax(scores, axis=-1, where=mask)
check(np.allclose(aw2, attention_weights, atol=1e-6))

mask_empty = mask.at[1].set(False)        # 第 1 行全部屏蔽
row_big = jax.nn.softmax(jnp.where(mask_empty, scores, -1e9), axis=-1)[1]
row_inf = jax.nn.softmax(jnp.where(mask_empty, scores, -jnp.inf), axis=-1)[1]
print("-1e9 →", row_big, "| -inf →", row_inf)
check(np.allclose(row_big, 0.25) and np.isnan(row_inf).all())
row_where = jax.nn.softmax(scores, axis=-1, where=mask_empty)[1]
print("where= →", row_where)
check(np.allclose(row_where, 0.0))          # where= 时全屏蔽行输出全 0

mask_float = mask.astype(padded_sequences.dtype)
mult_sum = jnp.sum(padded_sequences * mask_float, axis=-1)
mult_avg = mult_sum / jnp.maximum(jnp.sum(mask_float, axis=-1), 1e-9)
print("乘法求和:", mult_sum, "| 乘法平均:", mult_avg)
check(np.allclose(mult_sum, [6, 9, 30]) and np.allclose(mult_avg, [2, 4.5, 7.5]))

# 乘法掩码的陷阱：被屏蔽位置为 inf 时
v = jnp.array([1.0, jnp.inf]); m = jnp.array([1.0, 0.0])
print("v*m =", v * m, "| where(m>0, v, 0) =", jnp.where(m > 0, v, 0.0))
check(np.isnan((v * m)[1]) and float(jnp.where(m > 0, v, 0.0)[1]) == 0.0)

def masked_scan_body(hidden_state, inp):
    input_element, current_mask = inp                    # 掩码从 xs 中逐步读取
    potential_new_state = hidden_state * 0.9 + input_element * 0.1
    new_state = jnp.where(current_mask, potential_new_state, hidden_state)   # 无效步：保持不变
    output = jnp.where(current_mask, new_state * 2.0, 0.0)                   # 无效步：输出 0
    return new_state, output

@jax.jit
def masked_run(inputs, masks):
    # vmap 处理批维；scan 处理时间维
    return jax.vmap(lambda xs, ms: lax.scan(masked_scan_body, 0.0, (xs, ms)))(inputs, masks)

finals, outs = masked_run(padded_sequences, mask)
print("最终状态:", finals)

# 对拍：每条序列只在真实长度上运行 Python 循环（float64）
seqs = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
ref = []
for s in seqs:
    h = 0.0
    for x_t in s:
        h = h * 0.9 + x_t * 0.1
    ref.append(h)
check(np.allclose(finals, ref, atol=1e-6))
check(np.allclose(np.asarray(outs)[~np.asarray(mask)], 0.0))
print("填充 + 掩码的结果与逐条真实长度运行一致 ✓")

data = jnp.array([1., 2., 3., 4., 5., 6.])
seg = jnp.array([0, 0, 1, 1, 1, 2])
ssum = jax.ops.segment_sum(data, seg, num_segments=3)
smax = jax.ops.segment_max(data, seg, num_segments=4)        # 段 3 为空
print("segment_sum:", ssum, "| segment_max（含空段）:", smax)
check(np.allclose(ssum, [3, 12, 6]))
check(np.allclose(smax[:3], [2, 5, 6]) and np.isneginf(smax[3]))

# id 越界：被静默丢弃（6.0 不见了）
dropped = jax.ops.segment_sum(data, jnp.array([0, 0, 1, 1, 1, 5]), num_segments=3)
print("越界 id 被丢弃:", dropped)
check(np.allclose(dropped, [3, 12, 0]))

# jit 下不给 num_segments → 报错；用 static_argnames 或 functools.partial 固定它
try:
    jax.jit(lambda d, s: jax.ops.segment_sum(d, s))(data, seg); err = None
except jax.errors.ConcretizationTypeError as e:
    err = type(e).__name__
print("jit 下省略 num_segments:", err)
check(err == "ConcretizationTypeError")

seg_mean = jax.jit(lambda d, s: jax.ops.segment_sum(d, s, num_segments=3)
                   / jnp.maximum(jax.ops.segment_sum(jnp.ones_like(d), s, num_segments=3), 1))
print("分段平均:", seg_mean(data, seg))
check(np.allclose(seg_mean(data, seg), [1.5, 4.0, 6.0]))