第 3 章　用 grad 做自动微分

import jax
import jax.numpy as jnp
import numpy as np

print("JAX 版本:", jax.__version__)
print("默认后端:", jax.default_backend())

# 用 JAX 验证上面的两个例子

# 1) f(x) = x^2 在 x=2 处的导数（切线斜率）应为 4
print("d/dx x^2 在 x=2:", jax.grad(lambda x: x**2)(2.0))

# 2) f(x, y) = x^2 + sin(y) 在 (1, 0) 处的梯度应为 [2, 1]
def f_xy(v):
    x, y = v
    return x**2 + jnp.sin(y)

print("∇f(1, 0) =", jax.grad(f_xy)(jnp.array([1.0, 0.0])))

x0 = 1.0
exact = np.cos(x0)  # 真值（用 float64 计算）

print(f"{'h':>8} | {'前向差分误差':>14} | {'中心差分误差':>14}")
for h in [1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6]:
    x = jnp.float32(x0)
    fwd = (jnp.sin(x + h) - jnp.sin(x)) / h
    cen = (jnp.sin(x + h) - jnp.sin(x - h)) / (2 * h)
    print(f"{h:>8.0e} | {abs(float(fwd) - exact):>14.2e} | {abs(float(cen) - exact):>14.2e}")

print("\njax.grad 误差:", abs(float(jax.grad(jnp.sin)(jnp.float32(x0))) - exact))

import jax
import jax.numpy as jnp

def f(x):
    """计算 x 的三次方。"""
    return x**3

# 测试原函数
print(f"f(2.0) = {f(2.0)}")

# 使用 jax.grad 获得梯度函数
grad_f = jax.grad(f)

print(f"Type of f: {type(f)}")
print(f"Type of grad_f: {type(grad_f)}")

# 计算函数 f 在 x=2.0 处的梯度
gradient_at_2 = grad_f(2.0)

print(f"Gradient of f at x=2.0: {gradient_at_2}")  # 手算: 3 * 2.0**2 = 12.0

def g(x, y):
    """计算 x 的平方乘以 y。"""
    return (x**2) * y

# 获得函数 g 的梯度函数
grad_g = jax.grad(g)

# 计算在 (x=2.0, y=3.0) 处的梯度
# 默认情况下，这是相对于第一个参数 (x) 的梯度
gradient_g_wrt_x = grad_g(2.0, 3.0)

print(f"g(2.0, 3.0) = {g(2.0, 3.0)}")
print(f"Gradient of g w.r.t x at (2.0, 3.0): {gradient_g_wrt_x}")

# 验证：整数输入报错；allow_int 的行为；jacobian 即 jacrev
try:
    jax.grad(f)(2)
    int_input_raised = False
except TypeError as e:
    int_input_raised = True
    print("整数输入报错:", str(e).splitlines()[0])
assert int_input_raised, "整数输入应当报 TypeError"
assert jax.grad(f)(2.0) == 12.0

def scale(x, n):
    return x * n

# allow_int=True：整数参数可以出现在 argnums 中，梯度为 float0 类型的零
gx, gn = jax.grad(scale, argnums=(0, 1), allow_int=True)(2.0, 3)
print("allow_int=True ->", gx, "| n 的梯度是 float0 类型:", gn.dtype == jax.dtypes.float0)
assert gn.dtype == jax.dtypes.float0

# jax.jacobian 就是 jax.jacrev
import inspect
print("jax.jacobian 的文档:", jax.jacobian.__doc__.strip())
assert "jacrev(fun" in inspect.getsource(jax.jacobian)
F_demo = lambda x: jnp.stack([x[0] * x[1], jnp.sin(x[1])])
x_demo = jnp.array([1.0, 2.0])
assert jnp.array_equal(jax.jacobian(F_demo)(x_demo), jax.jacrev(F_demo)(x_demo))
print("断言通过：jacobian 与 jacrev 结果完全相同")

# ❌ 错误预期：以为 cos(pi/2) 的梯度会精确打印为 0.0
def h(x):
    """计算 sin(x)。"""
    return jnp.sin(x)

grad_h = jax.grad(h)
gradient_h_at_pi_half = grad_h(jnp.pi / 2.0)

print(f"h(pi/2) = {h(jnp.pi / 2.0)}")
print(f"Gradient of h at x=pi/2: {gradient_h_at_pi_half}")   # 数学值是 0.0
print("与 0.0 精确相等?", bool(gradient_h_at_pi_half == 0.0))

# ✅ 推荐写法：浮点结果用 allclose 比较
gradient_h_at_pi_half = grad_h(jnp.pi / 2.0)
assert gradient_h_at_pi_half != 0.0                          # 确实不是精确的 0
assert jnp.allclose(gradient_h_at_pi_half, 0.0, atol=1e-6)   # 但在 float32 精度内等于 0
print("断言通过：grad_h(pi/2) ≈ 0（atol=1e-6）")

# 用 JAX 验证上面手算的结果
def f_sin_sq(x):
    a = x**2         # 平方节点
    y = jnp.sin(a)   # sin 节点
    return y

x = 2.0
print("y      =", f_sin_sq(x))
print("ā = cos(4.0)      ≈", jnp.cos(4.0))
print("x̄ = ā * 2x        ≈", jnp.cos(4.0) * 2 * x)
print("jax.grad(f)(2.0)  =", jax.grad(f_sin_sq)(x))
assert jnp.allclose(jax.grad(f_sin_sq)(x), jnp.cos(4.0) * 2 * x)

print("原函数的 jaxpr：")
print(jax.make_jaxpr(f_sin_sq)(2.0))
print("梯度函数的 jaxpr：")
print(jax.make_jaxpr(jax.grad(f_sin_sq))(2.0))

def F(x):
    # F: R^3 -> R^2
    return jnp.array([x[0] * x[1], jnp.sin(x[2]) + x[0]])

x = jnp.array([1.0, 2.0, 3.0])

# 前向模式：J @ v，v 取第 1 个基向量 -> 得到雅可比第 1 列
y, Jv = jax.jvp(F, (x,), (jnp.array([1.0, 0.0, 0.0]),))
print("F(x) =", y, "| J 的第 1 列 =", Jv)

# 反向模式：u^T @ J，u 取第 1 个基向量 -> 得到雅可比第 1 行
y, vjp_fn = jax.vjp(F, x)
(uJ,) = vjp_fn(jnp.array([1.0, 0.0]))
print("J 的第 1 行 =", uJ)

print("完整雅可比 J =\n", jax.jacrev(F)(x))

# grad 就是输出维度为 1 时的 vjp
def loss(x):
    return jnp.sum(x**2)

_, vjp_loss = jax.vjp(loss, x)
print("\nvjp_fn(1.0) =", vjp_loss(1.0)[0], "| grad =", jax.grad(loss)(x))
J_full = jax.jacrev(F)(x)
assert jnp.allclose(Jv, J_full[:, 0]) and jnp.allclose(uJ, J_full[0, :])   # jvp 得列，vjp 得行
assert jnp.allclose(vjp_loss(1.0)[0], jax.grad(loss)(x))                    # grad 就是 vjp(1.0)

import time

def bench(fn, *args, n=20):
    """返回 fn(*args) 的平均耗时（毫秒）；先运行一次触发编译。"""
    jax.block_until_ready(fn(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        jax.block_until_ready(fn(*args))
    return (time.perf_counter() - t0) / n * 1e3

def layer(W, x):
    # 权重 W 作为参数传入；只对 x（argnums=1）求雅可比
    return jnp.tanh(W @ x)

W_wide = jax.random.normal(jax.random.key(0), (4, 2000))   # R^2000 -> R^4
W_tall = jax.random.normal(jax.random.key(1), (2000, 4))   # R^4 -> R^2000

jac_fwd = jax.jit(jax.jacfwd(layer, argnums=1))
jac_rev = jax.jit(jax.jacrev(layer, argnums=1))

results = {}
for name, W, x in [("R^2000 -> R^4 (n > m)", W_wide, jnp.ones(2000)),
                   ("R^4 -> R^2000 (n < m)", W_tall, jnp.ones(4))]:
    J_f, J_r = jac_fwd(W, x), jac_rev(W, x)
    # 正确性：jit 版本、非 jit 版本、两种模式之间都要一致
    assert jnp.allclose(J_f, J_r, atol=1e-5)
    assert jnp.allclose(J_r, jax.jacrev(layer, argnums=1)(W, x), atol=1e-5)
    t_f, t_r = bench(jac_fwd, W, x), bench(jac_rev, W, x)
    results[name] = (t_f, t_r)
    print(f"{name}: jacfwd {t_f:7.3f} ms | jacrev {t_r:7.3f} ms | 更快的是 {'jacfwd' if t_f < t_r else 'jacrev'}")

t_f, t_r = results["R^2000 -> R^4 (n > m)"]
print("\n规则「n > m 用 jacrev」是否得到印证:", t_r < t_f)
t_f, t_r = results["R^4 -> R^2000 (n < m)"]
print("规则「n < m 用 jacfwd」是否得到印证:", t_f < t_r)

# Hessian 的四种组合：结果相同，效率不同
def scalar_fn(W, x):
    return jnp.sum(jnp.tanh(W @ x) ** 2)

n = 300
W_sq = jax.random.normal(jax.random.key(2), (n, n)) / jnp.sqrt(n)
x0 = jax.random.normal(jax.random.key(3), (n,))

combos = {
    "jacfwd(jacrev)  [= jax.hessian]": jax.jacfwd(jax.jacrev(scalar_fn, argnums=1), argnums=1),
    "jacrev(jacrev)": jax.jacrev(jax.jacrev(scalar_fn, argnums=1), argnums=1),
    "jacrev(jacfwd)": jax.jacrev(jax.jacfwd(scalar_fn, argnums=1), argnums=1),
    "jacfwd(jacfwd)": jax.jacfwd(jax.jacfwd(scalar_fn, argnums=1), argnums=1),
}
H_ref = jax.hessian(scalar_fn, argnums=1)(W_sq, x0)
for name, fn in combos.items():
    fn_jit = jax.jit(fn)
    assert jnp.allclose(fn_jit(W_sq, x0), H_ref, atol=1e-4), name
    print(f"{name:34s}: {bench(fn_jit, W_sq, x0, n=10):7.2f} ms")
print("断言通过：四种组合得到的 Hessian 一致")

import jax
import jax.numpy as jnp

def power_product(x, y):
    """计算 x^2 * y"""
    return (x**2) * y

# 定义一些输入值
x_val = 3.0
y_val = 4.0

print(f"函数输出: {power_product(x_val, y_val)}")

# 对第一个参数 (x) 求导 - 默认行为
grad_f_wrt_x = jax.grad(power_product)
gradient_x = grad_f_wrt_x(x_val, y_val)

print(f"关于 x 的梯度 (默认): {gradient_x}")  # 预期: 2 * 3.0 * 4.0 = 24.0

# 使用 argnums=1 对第二个参数 (y) 求梯度
grad_f_wrt_y = jax.grad(power_product, argnums=1)
gradient_y = grad_f_wrt_y(x_val, y_val)

print(f"关于 y 的梯度 (argnums=1): {gradient_y}")  # 预期: 3.0**2 = 9.0

def affine(w, b, x):
    """计算 w*x + b"""
    return w * x + b

# 定义输入值
w_val = 2.0
b_val = 1.0
x_data = 5.0

print(f"仿射函数输出: {affine(w_val, b_val, x_data)}")

# 对第一个 (w) 和第二个 (b) 参数求梯度
grad_g_wrt_wb = jax.grad(affine, argnums=(0, 1))

# 注意: 在求导过程中，输入 x_data 被视为常数
gradients_wb = grad_g_wrt_wb(w_val, b_val, x_data)

print(f"关于 (w, b) 的梯度 (使用 argnums=(0, 1)): {gradients_wb}")
# 预期: (关于 w 的梯度, 关于 b 的梯度) = (x, 1) = (5.0, 1.0)

# 参数是一个 dict（pytree）
params = {
    "w": jnp.array([1.0, -2.0]),
    "b": jnp.array(0.5),
}
batch = {
    "inputs": jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),
    "targets": jnp.array([1.0, 0.0, -1.0]),
}

def model_apply(params, x):
    return x @ params["w"] + params["b"]

def loss_function(params, data_batch):
    predictions = model_apply(params, data_batch["inputs"])
    error = predictions - data_batch["targets"]
    return jnp.mean(error**2)  # 均方误差

grad_loss_wrt_params = jax.grad(loss_function, argnums=0)
gradients = grad_loss_wrt_params(params, batch)
print("gradients =", gradients)
assert jax.tree.structure(gradients) == jax.tree.structure(params)
print("断言通过：梯度的 pytree 结构与参数相同")

# ❌ 错误写法：参数是 dict 时不能直接做算术
learning_rate = 0.01
try:
    updated_params = params - learning_rate * gradients
    dict_arith_raised = False
except TypeError as e:
    dict_arith_raised = True
    print("报错 TypeError:", e)
assert dict_arith_raised

# ✅ 推荐写法：用 jax.tree.map 对每个叶子更新
updated_params = jax.tree.map(lambda p, g: p - learning_rate * g, params, gradients)

loss_before = loss_function(params, batch)
loss_after = loss_function(updated_params, batch)
print("更新前 loss:", loss_before)
print("更新后 loss:", loss_after)
assert loss_after < loss_before
assert jnp.allclose(updated_params["w"], params["w"] - learning_rate * gradients["w"])
print("断言通过：tree.map 正确更新了每个叶子，loss 下降")

def loss_with_metrics(params, data_batch):
    predictions = model_apply(params, data_batch["inputs"])
    loss = jnp.mean((predictions - data_batch["targets"]) ** 2)
    aux = {"predictions": predictions, "max_abs_err": jnp.max(jnp.abs(predictions - data_batch["targets"]))}
    return loss, aux

# 不加 has_aux：返回元组会报错
try:
    jax.grad(loss_with_metrics)(params, batch)
except TypeError as e:
    print("未加 has_aux 报错:", str(e).splitlines()[0])

# 加上 has_aux=True
grads, aux = jax.grad(loss_with_metrics, has_aux=True)(params, batch)
print("\ngrads:", grads)
print("aux  :", aux)

import jax
import jax.numpy as jnp

def f(x):
    return x**3

# 一阶导数
grad_f = jax.grad(f)
# 二阶导数：对一阶导数函数再求导
grad_grad_f = jax.grad(grad_f)

x = 2.0
print(f"f({x})   = {f(x)}")            # 8.0
print(f"f'({x})  = {grad_f(x)}")       # 3 * 2^2 = 12.0
print(f"f''({x}) = {grad_grad_f(x)}")  # 6 * 2 = 12.0

# 嵌套写法
second_derivative = jax.grad(jax.grad(f))
print("f''(2.0)  =", second_derivative(2.0))   # 12.0

# 三阶导数 f⁽³⁾(x) = 6
third_derivative = jax.grad(jax.grad(jax.grad(f)))
print("f⁽³⁾(2.0) =", third_derivative(2.0))    # 6.0

def f_xy(x, y):
    return x**2 * y + y**3

x, y = 2.0, 3.0

# ∂f/∂x
df_dx = jax.grad(f_xy, argnums=0)

# ∂²f/∂x²：再对 x 求导
d2f_dx2 = jax.grad(df_dx, argnums=0)
print("∂²f/∂x²  在 (2,3) =", d2f_dx2(x, y))   # 2y = 6.0

# ∂²f/∂y∂x：对 ∂f/∂x 再对 y 求导
d2f_dydx = jax.grad(df_dx, argnums=1)
print("∂²f/∂y∂x 在 (2,3) =", d2f_dydx(x, y))  # 2x = 4.0

# 把 (x, y) 打包成向量
def f_vec(v):
    x, y = v
    return x**2 * y + y**3

v = jnp.array([2.0, 3.0])
H = jax.hessian(f_vec)(v)
print("Hessian =\n", H)
# 解析解: [[2y, 2x], [2x, 6y]] = [[6, 4], [4, 18]]

assert jnp.allclose(H, jnp.array([[6.0, 4.0], [4.0, 18.0]]))
assert jnp.allclose(H, jax.jacfwd(jax.jacrev(f_vec))(v))
print("断言通过：与解析解、jacfwd(jacrev(f)) 一致")

# 多参数形式：argnums=(0, 1) 返回嵌套元组
H_tuple = jax.hessian(f_xy, argnums=(0, 1))(2.0, 3.0)
print("argnums=(0,1) 的结果:", H_tuple)

import jax
import jax.numpy as jnp

def polynomial(x):
    return x**3 + 2 * x**2 - 3 * x + 1

# 分开计算：前向计算会做两遍
value = polynomial(2.0)
gradient = jax.grad(polynomial)(2.0)
print(f"分开计算: value = {value}, gradient = {gradient}")

# 用 value_and_grad 一次得到两者
value_and_grad_fn = jax.value_and_grad(polynomial)
value, gradient = value_and_grad_fn(2.0)
print(f"value_and_grad: value = {value}, gradient = {gradient}")  # 11.0, 17.0

def multi_arg_func(a, b):
    return a**2 * jnp.sin(b)

a, b = 3.0, jnp.pi / 2

# 只对 a 求导
value, grad_a = jax.value_and_grad(multi_arg_func, argnums=0)(a, b)
print(f"argnums=0     : value = {value}, ∂/∂a = {grad_a}")  # 9.0, 6.0

# 对 a 和 b 都求导
value, (grad_a, grad_b) = jax.value_and_grad(multi_arg_func, argnums=(0, 1))(a, b)
print(f"argnums=(0,1) : value = {value}, ∂/∂a = {grad_a}, ∂/∂b = {grad_b}")

assert jnp.allclose(value, 9.0) and jnp.allclose(grad_a, 6.0)
assert grad_b != 0.0 and jnp.allclose(grad_b, 0.0, atol=1e-5)   # 不是精确 0，但在精度内为 0
print("断言通过：value=9, ∂/∂a=6, ∂/∂b≈0")

def loss_with_acc(w, x, y):
    logits = x * w
    loss = jnp.mean((logits - y) ** 2)
    acc = jnp.mean(jnp.sign(logits) == jnp.sign(y))   # 辅助信息：符号预测准确率
    return loss, acc

x = jnp.array([1.0, -2.0, 3.0])
y = jnp.array([2.0, -4.0, -1.0])

(loss, acc), g = jax.value_and_grad(loss_with_acc, has_aux=True)(1.0, x, y)
print(f"loss = {loss:.4f}, acc = {acc:.3f}, grad = {g:.4f}")
# 解析梯度：mean(2 (w x - y) x)
assert jnp.allclose(g, jnp.mean(2 * (1.0 * x - y) * x))

import jax
import jax.numpy as jnp

# ❌ 常见误解：以为下面对 x=-2.0 的调用会因「追踪器真值不明确」而报错
def conditional_func(x):
    if x > 0:
        return x * x
    else:
        return -x

grad_conditional = jax.grad(conditional_func)

try:
    g_pos, g_neg = grad_conditional(2.0), grad_conditional(-2.0)
    print("实测：没有报错。grad 在 x=2.0 :", g_pos, "| x=-2.0 :", g_neg)
except Exception as e:
    g_pos = g_neg = None
    print("报错:", type(e).__name__, e)

# 断言：单独用 grad 时 Python if 正常工作，两个分支的梯度都正确
assert g_pos == 4.0 and g_neg == -1.0

# 原因：单独 grad 时追踪器带有具体数值
seen = []
def peek(x):
    print("  grad 内部看到的 x:", x)
    seen.append(repr(x))
    return x * x

jax.grad(peek)(2.0)
assert "primal=2.0" in seen[0]   # 单独 grad：追踪器带着具体数值 2.0

# 真正会报错的是 jit：jit 的追踪器是抽象值
try:
    jax.jit(grad_conditional)(-2.0)
    jit_if_raised = False
except jax.errors.TracerBoolConversionError as e:
    jit_if_raised = True
    print("\njit(grad) 报错: TracerBoolConversionError")
    print(" ", str(e).splitlines()[0])
assert jit_if_raised, "jit 下 Python if 作用于追踪值应当报错"
print("断言通过：grad 单独使用不报错，套上 jit 才报错")

# ✅ 推荐写法：lax.cond，grad 与 jit(grad) 都正确
def conditional_func_lax(x):
    return jax.lax.cond(
        x > 0,              # 条件
        lambda x: x * x,    # 条件为真时执行
        lambda x: -x,       # 条件为假时执行
        x,                  # 操作数
    )

grad_cond_lax = jax.grad(conditional_func_lax)
grad_cond_lax_jit = jax.jit(grad_cond_lax)
for x in [2.0, -2.0]:
    print(f"x={x:5}: grad = {grad_cond_lax(x)}, jit(grad) = {grad_cond_lax_jit(x)}")

assert grad_cond_lax(2.0) == 4.0 and grad_cond_lax(-2.0) == -1.0
assert grad_cond_lax_jit(2.0) == 4.0 and grad_cond_lax_jit(-2.0) == -1.0
print("断言通过")

def fixed_loop(x, n=3):
    y = x
    for _ in range(n):   # n 在追踪时是已知的 Python 整数，循环会被展开
        y = y * 2
    return y

print("fixed_loop 梯度:", jax.grad(fixed_loop)(1.0))   # 8.0
assert jax.jit(jax.grad(fixed_loop))(1.0) == 8.0

def dynamic_loop_lax(x, n):
    # body_fun(i, val) -> 新的 val
    return jax.lax.fori_loop(0, n, lambda i, val: val * 2, x)

# 单独 grad：n=3 是具体整数，fori_loop 被转成 scan，可以反向求导
print("fori_loop 梯度 (n=3):", jax.grad(dynamic_loop_lax)(1.0, 3))   # 8.0
assert jax.grad(dynamic_loop_lax)(1.0, 3) == 8.0

# ❌ 常见误解：以为 while_loop / 动态 fori_loop 可以用来做可反向求导的动态循环
def double_until(x):
    # 状态 (值, 计数器)，循环 3 次
    return jax.lax.while_loop(lambda c: c[1] < 3, lambda c: (c[0] * 2, c[1] + 1), (x, 0))[0]

def raises_reverse_mode_error(fn, *args):
    try:
        fn(*args)
        return False
    except ValueError as e:
        print("  报错:", str(e).splitlines()[0])
        return True

print("grad 穿过 while_loop：")
assert raises_reverse_mode_error(jax.grad(double_until), 1.0)
print("jit 中 n 为追踪值的 fori_loop：")
assert raises_reverse_mode_error(jax.jit(jax.grad(dynamic_loop_lax)), 1.0, 3)

# while_loop 支持前向模式
primal, tangent = jax.jvp(double_until, (1.0,), (1.0,))
print("while_loop + jvp:", primal, tangent)
assert tangent == 8.0

# ✅ 推荐写法 1：把循环次数声明为静态参数
grad_static = jax.jit(jax.grad(dynamic_loop_lax), static_argnums=1)
print("static_argnums:", grad_static(1.0, 3))
assert grad_static(1.0, 3) == 8.0

# ✅ 推荐写法 2：用 lax.scan 写固定次数的循环
def double_scan(x, n=3):
    y, _ = jax.lax.scan(lambda c, _: (c * 2, None), x, None, length=n)
    return y

print("scan + jit(grad):", jax.jit(jax.grad(double_scan))(1.0))
assert jax.jit(jax.grad(double_scan))(1.0) == 8.0
print("断言通过")

import jax
import jax.numpy as jnp
import numpy as np

# 课程示例：先转整数再平方
def cast_and_square(x):
    x_int = x.astype(jnp.int32)          # 阶梯函数：信息被截断
    return x_int.astype(jnp.float32) ** 2

g_cast = jax.grad(cast_and_square)(2.7)
print("cast_and_square(2.7) 的梯度:", g_cast)
assert g_cast == 0.0

# 其他不连续函数和比较运算：梯度为 0
for name, fn in [("sign", jnp.sign), ("round", jnp.round), ("floor", jnp.floor), ("ceil", jnp.ceil),
                 ("(x>1).astype(float)", lambda x: (x > 1).astype(jnp.float32))]:
    g = jax.grad(fn)(2.7)
    print(f"grad({name})(2.7) = {g}")
    assert g == 0.0

# 输出本身是整数时，grad 直接报错
try:
    jax.grad(lambda x: x.astype(jnp.int32))(2.7)
except TypeError as e:
    print("\n整数输出报错:", str(e).splitlines()[0])

# ❌ 常见误解：以为 grad(abs)(0.0) 通常为 0.0，grad(sign)(0.0) 常为 NaN 或与平台相关
course_claim = {"grad(abs)(0.0)": 0.0, "grad(sign)(0.0)": "NaN 或与平台相关"}
measured = {"grad(abs)(0.0)": jax.grad(jnp.abs)(0.0), "grad(sign)(0.0)": jax.grad(jnp.sign)(0.0)}
for k in course_claim:
    print(f"{k:16s} 常见说法: {course_claim[k]!s:18s} 实际: {measured[k]}")

# ✅ 正确结论：尖点处的取值是确定的
assert jax.grad(jnp.abs)(0.0) == 1.0
assert jax.grad(jnp.sign)(0.0) == 0.0
assert jax.grad(jax.nn.relu)(0.0) == 0.0
assert jnp.isinf(jax.grad(jnp.sqrt)(0.0))
assert jnp.all(jnp.isnan(jax.grad(jnp.linalg.norm)(jnp.zeros(3))))
print("断言通过：abs→1.0，sign→0.0，relu→0.0，sqrt→inf，norm(0)→nan")

# 平滑近似：用 sigmoid(k x) 近似阶跃函数，梯度不再为零
k = 5.0
g_smooth = jax.grad(lambda x: jax.nn.sigmoid(k * x))(0.1)
print("grad(sigmoid(k x))(0.1) =", g_smooth)
assert g_smooth > 0

# 手动拆开 where(x>0, log(x), 0.) 在 x=0 处的反向传播
x = jnp.float32(0.0)
y_bar = 1.0                                  # 上游梯度
a_bar = jnp.where(x > 0, y_bar, 0.0)         # 分给 log 分支的梯度：0
dlog_dx = 1.0 / x                            # log'(0) = inf
print("ā =", a_bar, "| log'(0) =", dlog_dx, "| ā × log'(0) =", a_bar * dlog_dx)
assert jnp.isnan(a_bar * dlog_dx)            # 0 × inf = nan

# ❌ 朴素写法（单个 where）：前向值正确，但梯度是 NaN
naive = {
    "sqrt":   (lambda x: jnp.where(x > 0, jnp.sqrt(x), 0.0),                       0.0),
    "log":    (lambda x: jnp.where(x > 0, jnp.log(x), 0.0),                        0.0),
    "norm":   (lambda v: jnp.where(jnp.sum(v**2) == 0, 0.0, jnp.linalg.norm(v)),   jnp.zeros(3)),
    "除法":   (lambda y: jnp.where(y != 0, 2.0 / y, 0.0),                          0.0),
    "arcsin": (lambda x: jnp.where(jnp.abs(x) < 1, jnp.arcsin(x), jnp.sign(x) * jnp.pi / 2), 1.0),
}
for name, (fn, x_bad) in naive.items():
    val, g = jax.value_and_grad(fn)(x_bad)
    print(f"{name:6s} 危险点: 值 = {val:.4f}, 梯度 = {g}")
    assert jnp.all(jnp.isfinite(val)) and jnp.any(jnp.isnan(g))   # 值有限、梯度 NaN

# ✅ 推荐写法（双 where）：危险分支永远看不到危险输入
def safe_sqrt(x):
    xs = jnp.where(x > 0, x, 1.0)
    return jnp.where(x > 0, jnp.sqrt(xs), 0.0)

def safe_log(x):
    xs = jnp.where(x > 0, x, 1.0)
    return jnp.where(x > 0, jnp.log(xs), 0.0)

def safe_norm(v):
    s = jnp.sum(v**2)
    ss = jnp.where(s == 0, 1.0, s)
    return jnp.where(s == 0, 0.0, jnp.sqrt(ss))

def safe_div(y, a=2.0):
    ys = jnp.where(y != 0, y, 1.0)
    return jnp.where(y != 0, a / ys, 0.0)

def safe_arcsin(x):
    inside = jnp.abs(x) < 1
    xs = jnp.where(inside, x, 0.0)
    return jnp.where(inside, jnp.arcsin(xs), jnp.sign(x) * jnp.pi / 2)

safe = {
    # 名称: (函数, 危险点, 正常点, 正常点处的解析梯度)
    "sqrt":   (safe_sqrt,   0.0,          4.0,                  0.25),
    "log":    (safe_log,    0.0,          2.0,                  0.5),
    "norm":   (safe_norm,   jnp.zeros(3), jnp.array([3., 4., 0.]), jnp.array([0.6, 0.8, 0.0])),
    "除法":   (safe_div,    0.0,          4.0,                  -2.0 / 16),
    "arcsin": (safe_arcsin, 1.0,          0.5,                  1 / jnp.sqrt(1 - 0.25)),
}
for name, (fn, x_bad, x_ok, g_expected) in safe.items():
    g_bad = jax.grad(fn)(x_bad)
    g_ok = jax.grad(fn)(x_ok)
    print(f"{name:6s} 危险点梯度 = {g_bad} | 正常点梯度 = {g_ok}")
    assert jnp.all(jnp.isfinite(g_bad)), name                     # 危险点梯度有限
    assert jnp.allclose(g_ok, g_expected, atol=1e-6), name        # 正常点梯度不受影响
    assert jnp.allclose(fn(x_ok), naive[name][0](x_ok))           # 前向值与朴素写法一致
print("断言通过：5 个场景的危险点梯度都是有限值，正常点梯度正确")

from jax.scipy.special import logsumexp

g_log = jax.grad(jnp.log)(1e-20)
print("grad(log)(1e-20) =", g_log)
assert jnp.allclose(g_log, 1e20, rtol=1e-5)

# 朴素的 log-sum-exp 会溢出，梯度也随之变成 nan
def naive_lse(x):
    return jnp.log(jnp.sum(jnp.exp(x)))

x_big = jnp.array([1000.0, 1000.0])
print("\n朴素实现  值:", naive_lse(x_big), "| 梯度:", jax.grad(naive_lse)(x_big))
print("logsumexp 值:", logsumexp(x_big), "| 梯度:", jax.grad(logsumexp)(x_big))
assert jnp.all(jnp.isnan(jax.grad(naive_lse)(x_big)))
assert jnp.allclose(jax.grad(logsumexp)(x_big), 0.5)

def log1pexp_naive(x):
    return jnp.log1p(jnp.exp(x))

print("朴素实现 grad(100.0) =", jax.grad(log1pexp_naive)(100.0))   # nan
assert jnp.isnan(jax.grad(log1pexp_naive)(100.0))

@jax.custom_jvp
def log1pexp(x):
    return jnp.log1p(jnp.exp(x))

@log1pexp.defjvp
def log1pexp_jvp(primals, tangents):
    (x,), (x_dot,) = primals, tangents
    y = log1pexp(x)
    # 导数的稳定形式：sigmoid(x)
    return y, x_dot * jax.nn.sigmoid(x)

print("custom_jvp grad(100.0) =", jax.grad(log1pexp)(100.0))            # 1.0
print("custom_jvp grad(0.0)   =", jax.grad(log1pexp)(0.0))              # 0.5
print("二阶导也能用 (0.0)     =", jax.grad(jax.grad(log1pexp))(0.0))    # σ(0)(1-σ(0)) = 0.25
assert jax.grad(log1pexp)(100.0) == 1.0
assert jnp.allclose(jax.grad(jax.grad(log1pexp))(0.0), 0.25)

from jax.lax import stop_gradient

def f_normal(x):
    return x * x                    # 导数 2x

def f_stopped(x):
    return x * stop_gradient(x)     # 第二个 x 被视为常数，导数 = x

print("grad(x*x)(3.0)                 =", jax.grad(f_normal)(3.0))    # 6.0
print("grad(x*stop_gradient(x))(3.0)  =", jax.grad(f_stopped)(3.0))   # 3.0
assert jax.grad(f_normal)(3.0) == 6.0 and jax.grad(f_stopped)(3.0) == 3.0

# 实用技巧：straight-through estimator
# 前向用 round（不可导），反向假装是恒等函数（梯度为 1）
def ste_round(x):
    return x + stop_gradient(jnp.round(x) - x)

print("\nste_round(2.7) =", ste_round(2.7), "| 梯度 =", jax.grad(ste_round)(2.7))
assert ste_round(2.7) == 3.0 and jax.grad(ste_round)(2.7) == 1.0

# ❌ 常见误解：在被求导函数里调用 NumPy，以为它会被当作常量、梯度为 0
x = jnp.array([1.0, 2.0, 3.0])

try:
    g_np = jax.grad(lambda x: np.sin(x).sum())(x)
    np_raised = False
    print("np.sin 版本梯度:", g_np)
except jax.errors.TracerArrayConversionError as e:
    np_raised = True
    print("np.sin 作用于追踪器 -> TracerArrayConversionError")
assert np_raised, "NumPy 函数作用于 tracer 时应当报错"

# np.sum 碰巧转调 x.sum()，梯度并不是 0
g_npsum = jax.grad(lambda x: np.sum(x**2))(x)
print("np.sum 版本梯度:", g_npsum)
assert jnp.allclose(g_npsum, 2 * x)   # 不是 0，而是正确的 2x

# ✅ 推荐写法：被求导函数中一律使用 jnp；NumPy 只用来生成与输入无关的常量
noise = np.random.default_rng(0).random(3)            # NumPy 常量
g_jnp = jax.grad(lambda x: jnp.sum(jnp.sin(x) * noise))(x)
print("jnp 版本梯度:", g_jnp)
assert jnp.allclose(g_jnp, jnp.cos(x) * noise)

# 调试数值：jax.debug.print 在 jit 中也会每次执行
@jax.jit
def debug_loss(x):
    jax.debug.print("debug_loss 收到的 x = {}", x)
    return jnp.sum(x**2)

_ = jax.grad(debug_loss)(x)

def vector_output(x):
    return x**2   # 输出是向量，不是标量

x = jnp.array([1.0, 2.0, 3.0])
try:
    jax.grad(vector_output)(x)
    vec_raised = False
except TypeError as e:
    vec_raised = True
    print("TypeError:", str(e).splitlines()[0])
assert vec_raised

# 替代方案
J = jax.jacrev(vector_output)(x)
print("\njacrev 得到完整雅可比（对角阵 2x）:\n", J)
assert jnp.allclose(J, jnp.diag(2 * x))
# 如果只需要逐元素函数「对每个元素的导数」，对输出求和再 grad 就行
print("grad(sum(f)) =", jax.grad(lambda x: jnp.sum(vector_output(x)))(x))

import jax.test_util   # 必须显式导入
from jax.test_util import check_grads

def g(x):
    return jnp.sum(jnp.sin(x) * x**2)

x = jnp.array([0.5, 1.0, 2.0])

# 官方工具：通过时不输出，失败时抛 AssertionError
check_grads(g, (x,), order=2, modes=["fwd", "rev"])
print("check_grads: g 的一、二阶梯度检验通过")

# 故意写错导数（多乘了 2），check_grads 能发现
@jax.custom_jvp
def wrong_sin(x):
    return jnp.sin(x)

@wrong_sin.defjvp
def wrong_sin_jvp(primals, tangents):
    (x,), (t,) = primals, tangents
    return jnp.sin(x), 2 * jnp.cos(x) * t   # 错误：正确应为 cos(x) * t

try:
    check_grads(wrong_sin, (x,), order=1, modes=["rev"])
    caught = False
except AssertionError as e:
    caught = True
    print("check_grads 发现错误梯度:", str(e).strip().splitlines()[0])
assert caught

# 手写中心差分
def numerical_grad(f, x, h=1e-3):
    grads = []
    for i in range(x.size):
        e = jnp.zeros_like(x).at[i].set(h)
        grads.append((f(x + e) - f(x - e)) / (2 * h))
    return jnp.array(grads)

print("\n手写中心差分:", numerical_grad(g, x))
print("jax.grad    :", jax.grad(g)(x))
assert jnp.allclose(numerical_grad(g, x), jax.grad(g)(x), rtol=1e-3)

# ❌ 旧写法（过时）：jax.config 子模块在 JAX 0.12 中已移除
try:
    from jax.config import config
    config.update("jax_enable_x64", True)
    old_import_ok = True
except ModuleNotFoundError as e:
    old_import_ok = False
    print("实测报错 ModuleNotFoundError:", e)
assert not old_import_ok

# ✅ 推荐写法
import jax
import jax.numpy as jnp
import numpy as np

# jax.config.update("jax_enable_x64", True)   # 需要 64 位时取消注释，并放在程序最开头
print("当前 jax_enable_x64 =", jax.config.jax_enable_x64)
assert hasattr(jax.config, "update")

print(f"JAX 版本: {jax.__version__}")
print(f"默认后端: {jax.default_backend()}")

def poly_func(x):
    return 3 * x**2 + 2 * x + 5

# 得到梯度函数
grad_poly = jax.grad(poly_func)

x_val = 4.0
gradient_val = grad_poly(x_val)
analytical = 6 * x_val + 2

print(f"f({x_val}) = {poly_func(x_val)}")
print(f"JAX 梯度: {gradient_val}")
print(f"解析梯度: {analytical}")
assert gradient_val == analytical == 26.0

def multi_var_func(x, y):
    return x**3 * y + 2 * x**2

x_val, y_val = 2.0, 3.0

# 默认对第一个参数 x 求导（argnums=0）
grad_x = jax.grad(multi_var_func)(x_val, y_val)
print(f"∂g/∂x: {grad_x}  (解析: {3 * x_val**2 * y_val + 4 * x_val})")

# argnums=1 对 y 求导
grad_y = jax.grad(multi_var_func, argnums=1)(x_val, y_val)
print(f"∂g/∂y: {grad_y}  (解析: {x_val**3})")

# argnums=(0, 1) 返回两个梯度组成的元组
grad_xy = jax.grad(multi_var_func, argnums=(0, 1))(x_val, y_val)
print(f"(∂g/∂x, ∂g/∂y): {grad_xy}")
assert grad_xy == (44.0, 8.0)