第 4 章　进阶自动微分

# 公共导入
import time
import inspect
from functools import partial

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.test_util import check_grads
from jax.ad_checkpoint import print_saved_residuals   # 查看反向传播保存了哪些残差（4.10 节使用）

print("JAX 版本:", jax.__version__)
print("设备:", jax.devices())

def f(x):
    return jnp.sin(x) * x**2

x = jnp.arange(3.0)
v = jnp.ones(3)
u = jnp.array([1.0, 2.0, 3.0])

# 第 1 步：线性化，得到 y = f(x) 和线性函数 f_lin: v -> J v
y, f_lin = jax.linearize(f, x)
print("f_lin 的 jaxpr（a~d 是预先算好的 residuals，如 sin x、cos x、x²；e 是切向量，只剩对 e 的乘加）：")
print(jax.make_jaxpr(f_lin)(v))

# f_lin(v) 与 jvp 一致
jvp_out = jax.jvp(f, (x,), (v,))[1]
assert jnp.allclose(f_lin(v), jvp_out)

# 第 2 步：转置线性函数，得到 u -> J^T u
f_T = jax.linear_transpose(f_lin, x)
(ct_from_transpose,) = f_T(u)

# 与 jax.vjp 完全一致：反向模式 = 线性化 + 转置
_, f_vjp = jax.vjp(f, x)
(ct_from_vjp,) = f_vjp(u)
print("linear_transpose(linearize):", ct_from_transpose)
print("vjp                        :", ct_from_vjp)
assert jnp.allclose(ct_from_transpose, ct_from_vjp)

# 与显式 Jacobian 对拍：J^T u
J = jax.jacfwd(f)(x)
assert jnp.allclose(J.T @ u, ct_from_vjp)

# linear_transpose 只接受线性函数：对 sin 这种非线性函数做转置会报错
try:
    jax.linear_transpose(jnp.sin, x)(x)
    raised = False
except NotImplementedError as e:
    raised = True
    print("NotImplementedError:", str(e)[:90])
assert raised

# linearize 的好处：residuals 只算一次，同一点上的多个方向可以用 vmap 批量求 J v
V = jnp.eye(3)
cols = jax.vmap(f_lin)(V)          # 每行是 J e_i，即 J 的第 i 列
assert jnp.allclose(cols.T, J)
print("vmap(f_lin) 得到的 Jacobian：\n", cols.T)

# 课程代码：简单例子
def f(x):
    return jnp.array([x[0]**2, x[0] * x[1]])

x_primal = jnp.array([2.0, 3.0])
v_tangent = jnp.array([1.0, 0.5])

primal_out, tangent_out = jax.jvp(f, (x_primal,), (v_tangent,))
print(f"Primal input (x): {x_primal}")
print(f"Tangent vector (v): {v_tangent}")
print(f"Primal output f(x): {primal_out}")
print(f"Tangent output (J(x)v): {tangent_out}")

# 手工推导：J = [[2*x0, 0], [x1, x0]] = [[4, 0], [3, 2]]，J v = [4, 4]
J_manual = jnp.array([[4.0, 0.0], [3.0, 2.0]])
assert jnp.allclose(primal_out, jnp.array([4.0, 6.0]))
assert jnp.allclose(tangent_out, J_manual @ v_tangent)

# ❌ 旧写法（过时/错误）：用 stop_gradient(inputs) 当“零切向量”，或传 None
def predict(params, inputs):
    return jnp.dot(inputs, params['w']) + params['b']

params = {'w': jnp.array([[1.0, 2.0], [3.0, 4.0]]), 'b': jnp.array([0.1, -0.1])}
inputs = jnp.array([10.0, 20.0])
tangents = {'w': jnp.ones_like(params['w']), 'b': jnp.zeros_like(params['b'])}

_, t_wrong = jax.jvp(predict, (params, inputs), (tangents, jax.lax.stop_gradient(inputs)))
print("stop_gradient(inputs) 作切向量 ->", t_wrong, "（多出了 inputs 方向的贡献 inputs @ w）")
try:
    jax.jvp(predict, (params, inputs), (tangents, None))
except TypeError as e:
    print("传 None ->", type(e).__name__, ":", str(e)[:80], "...")

# ✅ 推荐写法：不扰动的参数传零切向量（或用闭包固定）
primal_out, tangent_out = jax.jvp(predict, (params, inputs), (tangents, jnp.zeros_like(inputs)))
print("Inputs:", inputs)
print("Params:", params)
print("Tangents (perturbation):", tangents)
print("Primal output (prediction):", primal_out)
print("Tangent output (change in prediction due to perturbation):", tangent_out)

# 另一种写法：用闭包只对 params 做 jvp
_, tangent_out2 = jax.jvp(lambda p: predict(p, inputs), (params,), (tangents,))

assert jnp.allclose(primal_out, jnp.array([70.1, 99.9]))      # 第二项是 99.9，不是 100.9
assert jnp.allclose(tangent_out, jnp.array([30.0, 30.0]))
assert jnp.allclose(tangent_out2, tangent_out)
assert jnp.allclose(t_wrong, jnp.array([100.0, 130.0]))       # stop_gradient 不等于零切向量

# 切向量 dtype 与 primal 不一致 -> TypeError
try:
    jax.jvp(jnp.sin, (jnp.float32(1.0),), (1,))
    raised = False
except TypeError as e:
    raised = True
    print("TypeError:", str(e)[:110], "...")
assert raised

# 整数 primal 必须配 float0 切向量
def scale_by_int(x, n):
    return x * n

n = jnp.int32(3)
n_tangent = np.zeros((), dtype=jax.dtypes.float0)
_, t = jax.jvp(scale_by_int, (2.0, n), (1.0, n_tangent))
print("对 x 方向的导数:", t)
assert t == 3.0

# 课程代码：标量输出
def f(x, y):
    return x**2 * jnp.sin(y)

x_primal = 3.0
y_primal = jnp.pi / 2.0

y_out, vjp_fun = jax.vjp(f, x_primal, y_primal)
print("前向输出 f(x, y):", y_out)

cotangent_v = 1.0
primal_grads = vjp_fun(cotangent_v)
print("对 x 的梯度:", primal_grads[0])   # 预期 2*x*sin(y) = 6.0
print("对 y 的梯度:", primal_grads[1])   # 预期 x^2*cos(y) = 0.0

grad_f = jax.grad(f, argnums=(0, 1))
grads_via_grad = grad_f(x_primal, y_primal)
print("jax.grad 结果:", grads_via_grad)

assert jnp.allclose(primal_grads[0], 6.0)
assert jnp.allclose(primal_grads[1], 0.0, atol=1e-5)
assert all(jnp.allclose(a, b) for a, b in zip(primal_grads, grads_via_grad))

# 课程代码：非标量输出
def g(x):
    return jnp.array([x[0]**2, x[1]**3])

x_primal = jnp.array([2.0, 3.0])
y_out, vjp_fun = jax.vjp(g, x_primal)
print("y_out:", y_out)                       # [4., 27.]

cotangent_v = jnp.array([10.0, 20.0])
primal_grad = vjp_fun(cotangent_v)
print("u^T J:", primal_grad)                 # ([40., 540.],)

assert isinstance(primal_grad, tuple) and len(primal_grad) == 1
assert jnp.allclose(primal_grad[0], jnp.array([40.0, 540.0]))

y, vjp_sq = jax.vjp(lambda x: x**2, jnp.arange(1.0, 4.0))

# 同一个 vjp_fun，vmap 到 3 个基向量上 -> Jacobian 的 3 行
rows = jax.vmap(vjp_sq)(jnp.eye(3))[0]
assert jnp.allclose(rows, jnp.diag(2 * jnp.arange(1.0, 4.0)))
print("vmap(vjp_fun)(I) =\n", rows)

# 余切形状不匹配（不会自动广播）
for bad_ct in [1.0, jnp.ones(2)]:
    try:
        vjp_sq(bad_ct)
        raised = False
    except ValueError as e:
        raised = True
        print("ValueError:", str(e)[:100], "...")
    assert raised

# 课程代码 A：连续两次 grad
def f(x):
    return x**3 + 2 * x**2 - 3 * x + 1

grad_f = jax.grad(f)
grad_grad_f = jax.grad(grad_f)

x_val = 2.0
first_deriv = grad_f(x_val)
second_deriv = grad_grad_f(x_val)
print("f(x) = x^3 + 2x^2 - 3x + 1")
print(f"f'(x) = 3x^2 + 4x - 3,  f'({x_val}) = {first_deriv}")
print(f"f''(x) = 6x + 4,        f''({x_val}) = {second_deriv}")

assert first_deriv == 17.0 and second_deriv == 16.0

# 课程代码 B / C：两种 HVP
def hvp_forward_over_reverse(f, primals, tangents):
    # jvp 返回 (grad_f(x), H v)，取第二个输出
    return jax.jvp(jax.grad(f), primals, tangents)[1]

def hvp_reverse_over_reverse(f, primals, tangents):
    x, = primals
    v, = tangents
    # 用 lambda 捕获 v，对 grad(f)(x)·v 再求一次梯度
    return jax.grad(lambda x_lambda: jnp.dot(jax.grad(f)(x_lambda), v))(x)

# 课程代码 D：f(x, y) = x^2 y + y^3
def func(x):
    return x[0]**2 * x[1] + x[1]**3

x_primal = jnp.array([1.0, 2.0])
v_tangent = jnp.array([1.0, 0.0])
hvp_val = jax.jvp(jax.grad(func), (x_primal,), (v_tangent,))[1]
# 梯度 [2xy, x^2+3y^2]；Hessian [[2y, 2x], [2x, 6y]]，在 (1,2) 处 = [[4,2],[2,12]]，Hv = [4, 2]
print("HVP (fwd-over-rev):", hvp_val)
print("HVP (rev-over-rev):", hvp_reverse_over_reverse(func, (x_primal,), (v_tangent,)))

assert jnp.allclose(hvp_val, jnp.array([4.0, 2.0]))
assert jnp.allclose(hvp_forward_over_reverse(func, (x_primal,), (v_tangent,)),
                    hvp_reverse_over_reverse(func, (x_primal,), (v_tangent,)))

# 课程代码 E：三种方式求完整 Hessian
hessian_matrix = jax.hessian(func)(x_primal)
hessian_jacfwd = jax.jacfwd(jax.grad(func))(x_primal)
hessian_jacrev = jax.jacrev(jax.grad(func))(x_primal)
print("jax.hessian:\n", hessian_matrix)

H_expected = jnp.array([[4.0, 2.0], [2.0, 12.0]])
for H in (hessian_matrix, hessian_jacfwd, hessian_jacrev):
    assert jnp.allclose(H, H_expected)

# 课程代码 F：三阶导数
grad_grad_grad_f = jax.grad(grad_grad_f)
third_deriv = grad_grad_grad_f(x_val)
print("f^(3)(2) =", third_deriv)
assert third_deriv == 6.0

class Dual:
    """朴素对偶数：所有嵌套层共用同一个无穷小 ε（没有层级标签）"""
    def __init__(self, p, t):
        self.p, self.t = p, t

    @staticmethod
    def lift(o):
        return o if isinstance(o, Dual) else Dual(o, 0.0)

    def __add__(self, o):
        o = Dual.lift(o)
        return Dual(self.p + o.p, self.t + o.t)
    __radd__ = __add__

    def __mul__(self, o):
        o = Dual.lift(o)
        return Dual(self.p * o.p, self.p * o.t + self.t * o.p)
    __rmul__ = __mul__

def naive_D(f):
    def df(x):
        out = f(Dual(x, 1.0))
        return out.t if isinstance(out, Dual) else 0.0
    return df

# d/dx [ x * d/dy (x + y)|_{y=1} ] 在 x=2 处，正确答案是 1
naive = naive_D(lambda x: x * naive_D(lambda y: x + y)(1.0))(2.0)

def jax_D(f):
    return lambda x: jax.jvp(f, (x,), (1.0,))[1]

jax_result = jax_D(lambda x: x * jax_D(lambda y: x + y)(1.0))(2.0)
print("朴素对偶数:", naive, "（扰动混淆，错误）")
print("JAX jvp   :", jax_result, "（正确）")
assert naive == 2.0
assert jax_result == 1.0

# 用 grad 嵌套同样正确：d/dx [ x * d/dy (x*y) ] = d/dx [x^2] = 2x = 4
assert jax.grad(lambda x: x * jax.grad(lambda y: x * y)(3.0))(2.0) == 4.0

from jax.experimental import jet

# 约定验证：x^3 在 x=2 处的 1/2/3 阶导数是 12, 12, 6
y0, (d1, d2, d3) = jet.jet(lambda x: x**3, (2.0,), ((1.0, 0.0, 0.0),))
print("x^3 的 jet:", y0, d1, d2, d3)
assert (y0, d1, d2, d3) == (8.0, 12.0, 12.0, 6.0)

# 与嵌套 grad 对拍：f(x) = exp(sin x) * cos(x^2)，求 1~6 阶导数
def g(x):
    return jnp.exp(jnp.sin(x)) * jnp.cos(x**2)

K = 6
def jet_derivs(x):
    _, ys = jet.jet(g, (x,), ((1.0,) + (0.0,) * (K - 1),))
    return jnp.stack(ys)

def nested_grads(x):
    out, h = [], g
    for _ in range(K):
        h = jax.grad(h)
        out.append(h(x))
    return jnp.stack(out)

dj, dn = jet_derivs(0.3), nested_grads(0.3)
print("jet       :", dj)
print("嵌套 grad :", dn)
assert jnp.allclose(dj, dn, rtol=1e-4)

# 计算图大小对比：第 K 阶导数的 jaxpr 长度（字符数）
def kth(h, k):
    for _ in range(k):
        h = jax.grad(h)
    return h
len_nested = len(str(jax.make_jaxpr(kth(g, K))(0.3)))
len_jet = len(str(jax.make_jaxpr(jet_derivs)(0.3)))
print(f"{K} 阶导数 jaxpr 字符数：嵌套 grad = {len_nested}，jet（同时给出 1~{K} 阶）= {len_jet}")

# jet 不支持的原语（实测）
unsupported = {
    "lax.cond": lambda x: lax.cond(x > 0, jnp.sin, jnp.cos, x),
    "jnp.sort": lambda x: jnp.sort(jnp.stack([x, 2 * x])).sum(),
    "lax.while_loop": lambda x: lax.while_loop(lambda c: c[0] < 3, lambda c: (c[0] + 1, c[1] * x), (0, 1.0))[1],
    "jax.nn.relu(custom_jvp)": jax.nn.relu,
}
n_fail = 0
for name, fn in unsupported.items():
    try:
        jet.jet(fn, (0.5,), ((1.0, 0.0),))
        print(f"{name:26s} 支持")
    except Exception as e:
        n_fail += 1
        print(f"{name:26s} 不支持 -> {type(e).__name__}")
assert n_fail == len(unsupported)

# 课程代码 A / B：jacfwd 与 jacrev
def func(x):
    return jnp.array([x[0]**2 * x[1], jnp.sin(x[2])])

x_in = jnp.array([1.0, 2.0, jnp.pi / 2])
jacobian_fwd = jax.jacfwd(func)(x_in)
jacobian_rev = jax.jacrev(func)(x_in)
print("输入:", x_in, " 输出:", func(x_in))
print("jacfwd (形状 %s):\n" % (jacobian_fwd.shape,), jacobian_fwd)

J_expected = jnp.array([[4.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
assert jacobian_fwd.shape == (2, 3)
assert jnp.allclose(jacobian_fwd, J_expected, atol=1e-6)
assert jnp.allclose(jacobian_rev, jacobian_fwd)

# 课程代码 C：vmap + jvp 手动构造（每个输入基向量 -> 一列）
n = x_in.shape[0]
basis_vectors_in = jnp.eye(n)
primals_out, jac_cols = jax.vmap(lambda tangent: jax.jvp(func, (x_in,), (tangent,)),
                                 in_axes=0)(basis_vectors_in)
jacobian_vmap_jvp = jac_cols.T      # vmap 把列堆在第 0 维，需转置

# 课程代码 D：vmap + vjp 手动构造（每个输出基向量 -> 一行）
primals_out, vjp_fun = jax.vjp(func, x_in)
m = primals_out.shape[0]
basis_vectors_out = jnp.eye(m)
jac_rows = jax.vmap(vjp_fun, in_axes=0)(basis_vectors_out)
jacobian_vmap_vjp = jac_rows[0]     # vjp_fun 返回一个元组，取第一个元素

assert jnp.allclose(jacobian_vmap_jvp, jacobian_fwd)
assert jnp.allclose(jacobian_vmap_vjp, jacobian_rev)
print("vmap(jvp) 与 vmap(vjp) 构造的 Jacobian 与 jacfwd/jacrev 一致")

# 查看 jax.hessian 的源码
src = inspect.getsource(jax.hessian)
print(src[src.index("return jacfwd"):].strip())
assert "jacfwd(jacrev(" in src

# 课程代码 E：jax.hessian
def scalar_func(x):
    return x[0]**2 * x[1] + x[1]**3

x_in = jnp.array([1.0, 2.0])
hessian_matrix = jax.hessian(scalar_func)(x_in)
print("Hessian:\n", hessian_matrix)
assert hessian_matrix.shape == (2, 2)
assert jnp.allclose(hessian_matrix, jnp.array([[4.0, 2.0], [2.0, 12.0]]))

# ❌ 旧写法（过时/错误）：切向量写成 (v_)，变量未定义且不是元组
v = jnp.array([0.5, -0.5])
grad_f = jax.grad(scalar_func)
try:
    _, hvp_efficient = jax.jvp(grad_f, (x_in,), (v_))
except NameError as e:
    print("NameError:", e)

# ✅ 推荐写法
v = jnp.array([0.5, -0.5])
# 方法 1：显式 Hessian
hessian_matrix = jax.hessian(scalar_func)(x_in)
hvp_explicit = hessian_matrix @ v
# 方法 2：jvp(grad) 直接得到 H v，不构造 H
grad_f = jax.grad(scalar_func)
_, hvp_efficient = jax.jvp(grad_f, (x_in,), (v,))
print("显式 H @ v :", hvp_explicit)
print("jvp(grad)  :", hvp_efficient)
assert jnp.allclose(hvp_explicit, jnp.array([1.0, -5.0]))
assert jnp.allclose(hvp_efficient, hvp_explicit)

def f_band(x):
    """f_i = x_i^2 - x_{i-1} x_{i+1} + sin(x_i)，Jacobian 为三对角"""
    xp = jnp.concatenate([x[1:], jnp.zeros(1, x.dtype)])
    xm = jnp.concatenate([jnp.zeros(1, x.dtype), x[:-1]])
    return x**2 - xm * xp + jnp.sin(x)

def colored_jacobian_tridiag(f, x, n_colors=3):
    n = x.shape[0]
    j = jnp.arange(n)
    seeds = (j[None, :] % n_colors == jnp.arange(n_colors)[:, None]).astype(x.dtype)  # (3, n)
    compressed = jax.vmap(lambda s: jax.jvp(f, (x,), (s,))[1])(seeds)                 # (3, n)：只做 3 次 JVP
    ii, jj = j[:, None], j[None, :]
    return jnp.where(jnp.abs(ii - jj) <= 1, compressed[jj % n_colors, ii], 0.0)

x9 = jnp.linspace(0.1, 1.0, 9)
J_dense = jax.jacfwd(f_band)(x9)
J_color = colored_jacobian_tridiag(f_band, x9)
print("三对角 Jacobian（着色法，3 次 JVP）:\n", np.round(np.asarray(J_color), 3))
assert jnp.allclose(J_color, J_dense)

# 计时对比（n = 3000）：jacfwd 需要 3000 次 JVP 并输出 n×n 稠密矩阵；
# 着色法只做 3 次 JVP，输出 3×n 的压缩形式（已包含全部非零元，三对角结构下无需解压成稠密）
def compressed_tridiag(f, x, n_colors=3):
    n = x.shape[0]
    j = jnp.arange(n)
    seeds = (j[None, :] % n_colors == jnp.arange(n_colors)[:, None]).astype(x.dtype)
    return jax.vmap(lambda s: jax.jvp(f, (x,), (s,))[1])(seeds)

n_big = 3000
x_big = jnp.linspace(0.1, 1.0, n_big)
jac_dense_fn = jax.jit(jax.jacfwd(f_band))
jac_comp_fn = jax.jit(lambda x: compressed_tridiag(f_band, x))

# 预热（编译）
jac_dense_fn(x_big).block_until_ready()
jac_comp_fn(x_big).block_until_ready()

def bench(fn, x, reps=10):
    t0 = time.perf_counter()
    for _ in range(reps):
        fn(x).block_until_ready()
    return (time.perf_counter() - t0) / reps * 1e3

t_dense, t_comp = bench(jac_dense_fn, x_big), bench(jac_comp_fn, x_big)
print(f"jacfwd（稠密 {n_big}×{n_big}）: {t_dense:.2f} ms")
print(f"着色法（压缩 3×{n_big}）    : {t_comp:.2f} ms   （本机比值约 {t_dense / t_comp:.1f}×）")

# 正确性：压缩形式中的对角线与稠密 Jacobian 一致
Jd, C = jac_dense_fn(x_big), jac_comp_fn(x_big)
idx = jnp.arange(n_big)
assert jnp.allclose(C[idx % 3, idx], jnp.diag(Jd))                      # 主对角 J[i,i]
assert jnp.allclose(C[(idx[1:]) % 3, idx[:-1]], jnp.diag(Jd, k=1))      # 上对角 J[i,i+1]
assert jnp.allclose(C[(idx[:-1]) % 3, idx[1:]], jnp.diag(Jd, k=-1))     # 下对角 J[i+1,i]

# ❌ 旧写法（过时/错误）：创建了从未使用的 jax.random.PRNGKey(0)（旧式 uint32[2] 原始 key）
key_old = jax.random.PRNGKey(0)
print("PRNGKey:", key_old.dtype, key_old.shape, "（旧写法，且本例根本不需要随机数）")

# ✅ 推荐写法：需要随机数时用 jax.random.key（类型化 key）；本例不需要，直接删掉
key_new = jax.random.key(0)
print("jax.random.key:", key_new.dtype, key_new.shape)
assert jnp.issubdtype(key_new.dtype, jax.dtypes.prng_key)

# 课程代码 B：f(x, y) = x * sin(y) 的自定义 VJP
@jax.custom_vjp
def f(x, y):
    return x * jnp.sin(y)

def f_fwd(x, y):
    primals_out = f(x, y)                         # 在 fwd 里调用 f 本身是允许的（此时不会再走自定义规则）
    residuals = (x, y, jnp.sin(y), jnp.cos(y))   # 这里保存了 4 个量；其实 (x, sin y, cos y) 就够了
    return primals_out, residuals

def f_bwd(residuals, cotangents_in):
    x, y, sin_y, cos_y = residuals
    g = cotangents_in
    grad_x = g * sin_y
    grad_y = g * x * cos_y
    return (grad_x, grad_y)

f.defvjp(f_fwd, f_bwd)

x_val = jnp.float32(2.0)
y_val = jnp.pi / 2.0
value, grads = jax.value_and_grad(lambda x, y: f(x, y)**2, argnums=(0, 1))(x_val, y_val)
value_auto, grads_auto = jax.value_and_grad(lambda x, y: (x * jnp.sin(y))**2, argnums=(0, 1))(x_val, y_val)
print("自定义:", value, grads)
print("自动  :", value_auto, grads_auto)
assert jnp.allclose(value, 4.0) and jnp.allclose(grads[0], 4.0) and jnp.allclose(grads[1], 0.0, atol=1e-5)
assert all(jnp.allclose(a, b, atol=1e-6) for a, b in zip(grads, grads_auto))

# ❌ 旧写法（过时/错误）：前向仍会溢出；try/except FloatingPointError 捕获不到任何东西
@jax.custom_vjp
def stable_log1pexp(x):
    return jnp.log1p(jnp.exp(x))

def stable_log1pexp_fwd(x):
    exp_x = jnp.exp(x)
    output = jnp.log1p(exp_x)
    return output, (x,)

def stable_log1pexp_bwd(residuals, cotangent_in):
    (x,) = residuals
    return (cotangent_in * (1.0 / (1.0 + jnp.exp(-x))),)

stable_log1pexp.defvjp(stable_log1pexp_fwd, stable_log1pexp_bwd)

def unstable_log1pexp(x):
    return jnp.log(1.0 + jnp.exp(x))

grad_stable = jax.grad(stable_log1pexp)
grad_unstable = jax.grad(unstable_log1pexp)
x_small, x_problematic, x_large = jnp.float32(1.0), jnp.float32(35.0), jnp.float32(100.0)
print(f"x=1   稳定梯度 {grad_stable(x_small):.6f}  不稳定梯度 {grad_unstable(x_small):.6f}")
print(f"x=35  稳定梯度 {grad_stable(x_problematic):.6f}  不稳定梯度 {grad_unstable(x_problematic):.6f}")
print(f"x=100 稳定梯度 {grad_stable(x_large)}  但前向 stable_log1pexp(100) = {stable_log1pexp(x_large)}")

caught = False
try:
    g_bad = grad_unstable(x_large)
    print("没有抛异常，不稳定梯度 =", g_bad)
except FloatingPointError:
    caught = True
print("except 分支是否执行:", caught)

# ✅ 推荐写法：前向用 logaddexp（稳定），反向用 sigmoid
@jax.custom_vjp
def log1pexp(x):
    return jnp.logaddexp(x, 0.0)

def log1pexp_fwd(x):
    return log1pexp(x), x

def log1pexp_bwd(x, g):
    return (g * jax.nn.sigmoid(x),)

log1pexp.defvjp(log1pexp_fwd, log1pexp_bwd)

print("log1pexp(100) =", log1pexp(100.0), " grad =", jax.grad(log1pexp)(100.0))
print("内置 jax.nn.softplus 的梯度 =", jax.grad(jax.nn.softplus)(100.0))

assert jnp.isinf(stable_log1pexp(x_large))           # ❌ 版本的前向溢出
assert log1pexp(100.0) == 100.0                      # 修正版前向正确
assert not caught and jnp.isnan(grad_unstable(x_large))   # 没有异常，静默返回 NaN
assert jax.grad(log1pexp)(100.0) == 1.0

# 想让 NaN 立刻报错，用 debug_nans（调试用，有性能代价）
with jax.debug_nans(True):
    try:
        grad_unstable(x_large)
        raised = False
    except FloatingPointError:
        raised = True
print("debug_nans 下是否抛 FloatingPointError:", raised)
assert raised

# bwd 对某个输入返回 None，JAX 视为零梯度
@jax.custom_vjp
def scaled(x, s):
    return x * s

scaled.defvjp(lambda x, s: (x * s, s), lambda s, g: (g * s, None))
gx, gs = jax.grad(scaled, argnums=(0, 1))(2.0, 3.0)
print("对 x 的梯度:", gx, " 对 s 的梯度（bwd 返回 None）:", gs)
assert gx == 3.0 and gs == 0.0

# 梯度裁剪：前向恒等，反向把余切裁剪到 [lo, hi]
@partial(jax.custom_vjp, nondiff_argnums=(0, 1))
def clip_gradient(lo, hi, x):
    return x

def clip_gradient_fwd(lo, hi, x):
    return x, None                        # 无需 residuals

def clip_gradient_bwd(lo, hi, res, g):    # 不可微参数排在最前面
    return (jnp.clip(g, lo, hi),)         # 只为可微参数 x 返回梯度

clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)

loss = lambda x: jnp.sum(10.0 * clip_gradient(-1.0, 1.0, x) ** 2)
xs = jnp.array([-2.0, 0.01, 0.03, 3.0])
print("原始梯度  :", jax.grad(lambda x: jnp.sum(10.0 * x**2))(xs))
print("裁剪后梯度:", jax.grad(loss)(xs))
assert jnp.allclose(jax.grad(loss)(xs), jnp.clip(20.0 * xs, -1.0, 1.0))
assert jnp.allclose(clip_gradient(-1.0, 1.0, xs), xs)   # 前向恒等

# 梯度反转层（Gradient Reversal Layer, DANN）：前向恒等，反向乘以 -λ
@jax.custom_vjp
def grad_reverse(x, lam):
    return x

def grad_reverse_fwd(x, lam):
    return x, lam

def grad_reverse_bwd(lam, g):
    return (-lam * g, None)        # λ 是数组参数（可被 jit 追踪），不求导，返回 None

grad_reverse.defvjp(grad_reverse_fwd, grad_reverse_bwd)

# 特征提取器 h = w*x 接一个“域分类器” loss = h^2；GRL 让 w 的梯度反号
def domain_loss(w, x, lam):
    h = w * x
    return grad_reverse(h, lam) ** 2

w0, x0, lam0 = 1.5, 2.0, 0.5
g_plain = jax.grad(lambda w: (w * x0) ** 2)(w0)
g_grl = jax.jit(jax.grad(domain_loss))(w0, x0, lam0)
print("普通梯度:", g_plain, "  GRL 梯度:", g_grl)
assert jnp.allclose(g_grl, -lam0 * g_plain)

# 陷阱：把 jit 追踪的数组放进 nondiff_argnums -> UnexpectedTracerError
@partial(jax.custom_vjp, nondiff_argnums=(1,))
def grl_bad(x, lam):
    return x
grl_bad.defvjp(lambda x, lam: (x, None), lambda lam, r, g: (-lam * g,))
print("lam 是 Python 常量时可用:", jax.grad(lambda x: grl_bad(x, 0.5))(1.0))
try:
    jax.jit(jax.grad(lambda x, lam: grl_bad(x, lam)))(1.0, 0.5)
    raised = False
except jax.errors.UnexpectedTracerError as e:
    raised = True
    print("UnexpectedTracerError:", str(e)[:100], "...")
assert raised

@partial(jax.custom_vjp, nondiff_argnums=(0,))
def fixed_point(T, a, x0):
    """用 while_loop 求 x = T(a, x) 的不动点（迭代次数由收敛条件决定）"""
    def cond(c):
        x, x_prev, i = c
        return (jnp.max(jnp.abs(x - x_prev)) > 1e-6) & (i < 500)
    def body(c):
        x, _, i = c
        return T(a, x), x, i + 1
    return lax.while_loop(cond, body, (T(a, x0), x0, 0))[0]

def fixed_point_fwd(T, a, x0):
    x_star = fixed_point(T, a, x0)
    return x_star, (a, x_star)                 # 只保存解本身

def fixed_point_bwd(T, res, g):
    a, x_star = res
    _, vjp_a = jax.vjp(lambda a: T(a, x_star), a)
    _, vjp_x = jax.vjp(lambda x: T(a, x), x_star)
    # 伴随不动点 u = g + (∂T/∂x)^T u，也用 while_loop 迭代
    def cond(c):
        u, u_prev, i = c
        return (jnp.max(jnp.abs(u - u_prev)) > 1e-7) & (i < 500)
    def body(c):
        u, _, i = c
        return g + vjp_x(u)[0], u, i + 1
    u = lax.while_loop(cond, body, (g + vjp_x(g)[0], g, 0))[0]
    (a_bar,) = vjp_a(u)
    return a_bar, jnp.zeros_like(x_star)       # 解与初值无关，x0 的梯度为 0

fixed_point.defvjp(fixed_point_fwd, fixed_point_bwd)

# 例 1：Newton 迭代求 sqrt(a)：x = (x + a/x)/2，解析导数 d sqrt(a)/da = 1/(2 sqrt(a))
newton_T = lambda a, x: 0.5 * (x + a / x)
a = 2.0
x_star = fixed_point(newton_T, a, 1.0)
g_implicit = jax.grad(lambda a: fixed_point(newton_T, a, 1.0))(a)

def unrolled_newton(a, n_iter=30):               # 展开版本：固定步数，普通反向传播
    return lax.fori_loop(0, n_iter, lambda i, x: newton_T(a, x), 1.0)
g_unrolled = jax.grad(unrolled_newton)(a)

print(f"x* = {x_star:.7f}  (sqrt(2) = {np.sqrt(2):.7f})")
print(f"隐函数梯度 {g_implicit:.7f} | 展开梯度 {g_unrolled:.7f} | 解析 {0.5 / np.sqrt(2):.7f}")
assert jnp.allclose(x_star, np.sqrt(2.0))
assert jnp.allclose(g_implicit, 0.5 / np.sqrt(2.0))
assert jnp.allclose(g_implicit, g_unrolled)

# 例 2：x = cos(a x)，解析：dx/da = -x sin(ax) / (1 + a sin(ax))
cos_T = lambda a, x: jnp.cos(a * x)
a = 1.0
x_star = fixed_point(cos_T, a, 0.5)
g_implicit = jax.jit(jax.grad(lambda a: fixed_point(cos_T, a, 0.5)))(a)   # 可以 jit
g_analytic = -x_star * jnp.sin(a * x_star) / (1 + a * jnp.sin(a * x_star))
print(f"x* = {x_star:.6f}, 隐函数梯度 = {g_implicit:.6f}, 解析 = {g_analytic:.6f}")
assert jnp.allclose(g_implicit, g_analytic, atol=1e-5)

# 对照：直接对 while_loop 版本做反向求导会失败（4.8 节详解）
def fixed_point_plain(a):
    return lax.while_loop(lambda c: jnp.abs(c[0] - c[1]) > 1e-6,
                          lambda c: (cos_T(a, c[0]), c[0]), (cos_T(a, 0.5), 0.5))[0]
try:
    jax.grad(fixed_point_plain)(a)
    raised = False
except ValueError as e:
    raised = True
    print("直接 grad(while_loop):", type(e).__name__, "-", str(e)[:75], "...")
assert raised

# 1) 正确的自定义规则：通过
check_grads(f, (jnp.float32(2.0), jnp.float32(0.7)), order=1, modes=["rev"])
check_grads(log1pexp, (jnp.float32(3.0),), order=1, modes=["rev"])
check_grads(lambda a: fixed_point(cos_T, a, 0.5), (jnp.float32(1.0),), order=1, modes=["rev"], atol=1e-2, rtol=1e-2)
print("f / log1pexp / fixed_point 通过 check_grads")

# 2) 故意写错的 bwd：被抓住
@jax.custom_vjp
def buggy_sin(x):
    return jnp.sin(x)
buggy_sin.defvjp(lambda x: (jnp.sin(x), x), lambda x, g: (g * jnp.sin(x),))   # 应为 cos
try:
    check_grads(buggy_sin, (jnp.float32(1.0),), order=1, modes=["rev"])
    caught_bug = False
except AssertionError:
    caught_bug = True
print("check_grads 发现错误的 bwd:", caught_bug)
assert caught_bug

# 3) 故意修改梯度的规则（GRL）同样“通不过”，说明规则在起作用
try:
    # 注意：参数要传 jnp 标量。本例若传 Python float 1.0，check_grads 内部调用 .reshape 时会报 AttributeError（实测）
    check_grads(lambda x: grad_reverse(x, 0.5), (jnp.float32(1.0),), order=1, modes=["rev"])
    grl_passes = True
except AssertionError:
    grl_passes = False
assert not grl_passes

# ❌ 旧写法（过时/错误）：前向也用 safe_x，导致 stable_sqrt(0) = 1
@jax.custom_jvp
def stable_sqrt_course(x):
    is_zero = (x == 0.0)
    safe_x = jnp.where(is_zero, 1.0, x)
    return jnp.sqrt(safe_x)

@stable_sqrt_course.defjvp
def stable_sqrt_course_jvp(primals, tangents):
    x, = primals
    t, = tangents
    primal_out = stable_sqrt_course(x)
    is_zero = (x == 0.0)
    safe_x = jnp.where(is_zero, 1.0, x)
    deriv = 0.5 * jax.lax.rsqrt(safe_x)
    tangent_out = jnp.where(is_zero, 0.0, deriv * t)
    return primal_out, tangent_out

print("原始版 jvp at x=0:", jax.jvp(stable_sqrt_course, (0.0,), (1.0,)))

# ✅ 推荐写法：前向用真 sqrt，safe_x 只用于导数公式
@jax.custom_jvp
def stable_sqrt(x):
    return jnp.sqrt(x)

@stable_sqrt.defjvp
def stable_sqrt_jvp(primals, tangents):
    x, = primals
    t, = tangents
    is_zero = (x == 0.0)
    safe_x = jnp.where(is_zero, 1.0, x)
    tangent_out = jnp.where(is_zero, 0.0, 0.5 * lax.rsqrt(safe_x) * t)   # 约定 x=0 处导数为 0
    return jnp.sqrt(x), tangent_out

p0, t0 = jax.jvp(stable_sqrt, (0.0,), (1.0,))
p4, t4 = jax.jvp(stable_sqrt, (4.0,), (1.0,))
print("x=0:", p0, t0, "   x=4:", p4, t4)       # x=4: 0.5 * 4**-0.5 = 0.25
assert stable_sqrt_course(0.0) == 1.0           # ❌ 版本的前向错误
assert p0 == 0.0 and t0 == 0.0
assert p4 == 2.0 and t4 == 0.25
# 对比：原生 jnp.sqrt 在 0 处梯度为 inf
assert jnp.isinf(jax.grad(jnp.sqrt)(0.0)) and jax.grad(stable_sqrt)(0.0) == 0.0

# 1) custom_jvp 自动获得反向模式（grad 通过转置得到）
assert jax.grad(stable_sqrt)(4.0) == 0.25
assert jnp.allclose(jax.jacrev(stable_sqrt)(4.0), jax.jacfwd(stable_sqrt)(4.0))

# 2) custom_jvp 对象没有 defvjp：两者不能同时定义
assert not hasattr(stable_sqrt, "defvjp")
print("custom_jvp 对象的方法:", [m for m in dir(stable_sqrt) if m.startswith("def")])

# 3) custom_vjp 不支持前向模式：jvp / jacfwd / hessian 都报 TypeError
n_err = 0
for name, call in [("jvp", lambda: jax.jvp(log1pexp, (1.0,), (1.0,))),
                   ("jacfwd", lambda: jax.jacfwd(log1pexp)(1.0)),
                   ("hessian", lambda: jax.hessian(log1pexp)(1.0))]:
    try:
        call()
    except TypeError as e:
        n_err += 1
        print(f"{name:8s} -> TypeError: {str(e)[:70]}")
assert n_err == 3

@jax.custom_jvp
def softplus(x):
    # 数值稳定的前向：max(x,0) + log1p(exp(-|x|))，不会溢出
    return jnp.maximum(x, 0.0) + jnp.log1p(jnp.exp(-jnp.abs(x)))

softplus.defjvps(lambda t, ans, x: t * jax.nn.sigmoid(x))

xs = jnp.array([-100.0, -10.0, 0.0, 10.0, 100.0])
print("softplus:", softplus(xs))
print("grad    :", jax.vmap(jax.grad(softplus))(xs))
assert jnp.allclose(softplus(xs), jax.nn.softplus(xs))
assert jnp.allclose(jax.vmap(jax.grad(softplus))(xs), jax.nn.sigmoid(xs))
assert jnp.all(jnp.isfinite(jax.vmap(jax.grad(jax.grad(softplus)))(xs)))   # 二阶导也能用

# 前向和反向都能用，并通过 check_grads（同时检查 fwd 和 rev、二阶）
check_grads(softplus, (jnp.float32(1.3),), order=2, modes=["fwd", "rev"])

# 用 ans 复用前向结果：sigmoid' = ans * (1 - ans)；多参数时可对某参数传 None
@jax.custom_jvp
def my_sigmoid(x):
    return 1.0 / (1.0 + jnp.exp(-x))
my_sigmoid.defjvps(lambda t, ans, x: t * ans * (1.0 - ans))

@jax.custom_jvp
def mul_no_grad_y(x, y):
    return x * y
mul_no_grad_y.defjvps(lambda tx, ans, x, y: tx * y, None)   # y 不求导

assert jnp.allclose(jax.grad(my_sigmoid)(0.3), jax.grad(jax.nn.sigmoid)(0.3))
assert jax.grad(mul_no_grad_y, argnums=(0, 1))(2.0, 3.0) == (3.0, 0.0)
print("defjvps 检查通过")

@jax.custom_jvp
def nonlinear_rule(x):
    return jnp.sin(x)

@nonlinear_rule.defjvp
def _nonlinear_rule_jvp(primals, tangents):
    x, = primals
    t, = tangents
    return jnp.sin(x), jnp.cos(x) * t**2        # ❗ 对切向量非线性

print("jvp 能跑:", jax.jvp(nonlinear_rule, (1.0,), (2.0,)))
try:
    jax.grad(nonlinear_rule)(1.0)
    raised = False
except NotImplementedError as e:
    raised = True
    print("grad -> NotImplementedError:", str(e)[:70])
assert raised

# 规则里对 primal 用 while_loop 完全没问题：切向量部分仍是线性的
@jax.custom_jvp
def sin_with_loop(x):
    return jnp.sin(x)

@sin_with_loop.defjvp
def _sin_with_loop_jvp(primals, tangents):
    x, = primals
    t, = tangents
    d = lax.while_loop(lambda c: c[0] < 3, lambda c: (c[0] + 1, c[1]), (0, jnp.cos(x)))[1]
    return jnp.sin(x), d * t

assert jnp.allclose(jax.grad(sin_with_loop)(1.0), jnp.cos(1.0))

# 切向量交给 pure_callback：数学上仍是线性的，jvp 能跑，但 grad 在转置时失败
@jax.custom_jvp
def tangent_via_callback(x):
    return jnp.sin(x)

@tangent_via_callback.defjvp
def _tangent_via_callback_jvp(primals, tangents):
    x, = primals
    t, = tangents
    t_cb = jax.pure_callback(lambda a: np.asarray(a), jax.ShapeDtypeStruct(t.shape, t.dtype), t)
    return jnp.sin(x), jnp.cos(x) * t_cb

x1 = jnp.float32(1.0)
assert jnp.allclose(jax.jvp(tangent_via_callback, (x1,), (x1,))[1], jnp.cos(1.0))
try:
    jax.grad(tangent_via_callback)(x1)
    raised = False
except ValueError as e:
    raised = True
    print("切向量经 pure_callback → grad:", str(e)[:60])
assert raised