第 4 章　用 vmap 自动向量化

# 本章通用导入与工具函数
import time
import jax
import jax.numpy as jnp
from jax import lax
import numpy as np

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

def bench(fn, *args, n=20):
    # 先预热一次（触发编译），再计时 n 次取平均，单位毫秒
    fn(*args).block_until_ready()
    t0 = time.perf_counter()
    for _ in range(n):
        out = fn(*args)
    out.block_until_ready()
    return (time.perf_counter() - t0) / n * 1e3

def jaxpr_str(closed_jaxpr):
    # 把 jaxpr 转成不带终端颜色码的纯文本，便于用 assert 检查其中的内容
    return closed_jaxpr.pretty_print(use_color=False)

# 课程代码 1：逐个处理一批图像（predict 只是示意，这里补一个占位实现，方便运行）
def predict(image):
    return jnp.mean(image)          # 占位：对单张图像返回一个分数

def predict_batch_loop(images):
    results = []
    for image in images:
        results.append(predict(image))
    return results                  # 也可以用 jnp.stack 堆成一个数组

images = jnp.ones((3, 8, 8))        # 3 张 8x8 的“图像”
print(predict_batch_loop(images))

# 课程代码 2：纯 Python 循环
numbers = [1, 2, 3, 4, 5]
results = []
for x in numbers:
    results.append(x + 10)
print(results)                      # [11, 12, 13, 14, 15]

# 课程代码 3：向量化写法（NumPy）
import numpy as np
numbers = np.array([1, 2, 3, 4, 5])
results = numbers + 10              # 对整个数组一次完成加法，没有显式 Python 循环
print(results)                      # [11 12 13 14 15]
assert results.tolist() == [11, 12, 13, 14, 15]

key = jax.random.key(0)
k1, k2 = jax.random.split(key)
xs = jax.random.normal(k1, (1000, 64))   # 1000 个样本，每个是 64 维向量
w = jax.random.normal(k2, (64,))

def dot_loop(xs, w):                     # Python 循环：每次迭代都要派发一次设备操作
    return jnp.stack([jnp.dot(x, w) for x in xs])

def dot_manual(xs, w):                   # 手动向量化：一次矩阵-向量乘
    return xs @ w

assert jnp.allclose(dot_loop(xs, w), dot_manual(xs, w), atol=1e-5)
print("Python 循环:  %.2f ms" % bench(dot_loop, xs, w, n=3))
print("手动向量化:  %.3f ms" % bench(dot_manual, xs, w))

def normalize(x):            # 单样本：x 形状 (3,)，返回单位向量
    return x / jnp.linalg.norm(x)

X = jnp.array([[3.0, 4.0, 0.0],
               [1.0, 0.0, 0.0]])    # 2 个样本，正确结果的每一行范数都应为 1

wrong = normalize(X)                # 直接传入整批：不报错，但结果错误
right = jax.vmap(normalize)(X)      # vmap：每个样本各自归一化
print("直接传入整批:\n", wrong, "\n 各行范数:", jnp.linalg.norm(wrong, axis=1))
print("vmap:\n", right, "\n 各行范数:", jnp.linalg.norm(right, axis=1))
assert not jnp.allclose(jnp.linalg.norm(wrong, axis=1), 1.0)     # 直接传整批：结果错误
assert jnp.allclose(jnp.linalg.norm(right, axis=1), 1.0)         # vmap：每行都是单位向量

import jax
import jax.numpy as jnp

# 处理单个向量的函数（期望输入是一维数组）
def sum_of_squares(vector):
    print(f"Running sum_of_squares for a vector of shape: {vector.shape}")
    return jnp.sum(vector**2)

single_vector = jnp.array([1., 2., 3.])
result_single = sum_of_squares(single_vector)
print(f"Result for single vector: {result_single}")   # 1^2 + 2^2 + 3^2 = 14.0

# 一批向量：4 行，每行是一个 3 维向量
batch_of_vectors = jnp.array([
    [1., 2., 3.],
    [4., 5., 6.],
    [7., 8., 9.],
    [0., 1., 0.],
])
print(f"Batch shape: {batch_of_vectors.shape}")      # (4, 3)

# 手动循环（仅作演示，不推荐在 JAX 中这样写）
manual_results = []
for i in range(batch_of_vectors.shape[0]):
    vector = batch_of_vectors[i]
    manual_results.append(sum_of_squares(vector))
manual_output = jnp.stack(manual_results)
print(f"Manual loop output: {manual_output}")        # 函数里的 print 会执行 4 次

# 使用 vmap 创建向量化版本
vectorized_sum_of_squares = jax.vmap(sum_of_squares)
vmap_output = vectorized_sum_of_squares(batch_of_vectors)
print(f"vmap output: {vmap_output}")                 # 注意：函数里的 print 只执行了 1 次

def show(v):
    print("Python print（追踪时执行）:", type(v).__name__, v.shape)
    jax.debug.print("jax.debug.print（运行时执行）: {}", v)
    return jnp.sum(v ** 2)

import io, contextlib
buf = io.StringIO()
with contextlib.redirect_stdout(buf):            # 捕获输出，用来验证打印次数
    res = jax.vmap(show)(batch_of_vectors[:3]).block_until_ready()
    jax.effects_barrier()                        # 等待 debug.print 的回调全部完成
captured = buf.getvalue()
print(captured, "结果:", res)

# 验证：Python print 只在追踪时执行 1 次；jax.debug.print 对 3 个样本各执行 1 次
assert captured.count("Python print") == 1
assert captured.count("jax.debug.print") == 3
print("✔ 断言通过：print 只打印一次是追踪造成的，并不是计算被“优化掉”了")

def sos(v):                               # 去掉 print 的 sum_of_squares
    return jnp.sum(v ** 2)

print("===== f 的 jaxpr =====")
print(jax.make_jaxpr(sos)(jnp.ones(3)))
print("===== vmap(f) 的 jaxpr =====")
jp_v = jaxpr_str(jax.make_jaxpr(jax.vmap(sos))(batch_of_vectors))
print(jp_v)
# 验证：vmap 后的程序里没有任何循环原语，只是形状多了一维
assert not any(p in jp_v for p in ("while", "scan", "fori")) and "f32[4]" in jp_v

def dot(x, w):
    return jnp.dot(x, w)
print("===== vmap(dot) 的 jaxpr：dot_general 多了批维度 ([0], [0]) =====")
print(jax.make_jaxpr(jax.vmap(dot))(jnp.ones((5, 3)), jnp.ones((5, 3))))

def relu_py(x):
    if x > 0:          # Python 控制流依赖了被映射的值
        return x
    return 0.0

try:
    jax.vmap(relu_py)(jnp.array([-1.0, 2.0, 3.0]))
    raised = False
except jax.errors.TracerBoolConversionError as e:
    raised = True
    print(type(e).__name__, ":", str(e).splitlines()[0])
assert raised, "依赖批数据的 Python if 应当报错"

# 正确写法：用 jnp.where 表达逐元素选择
relu = lambda x: jnp.where(x > 0, x, 0.0)
out = jax.vmap(relu)(jnp.array([-1.0, 2.0, 3.0]))
assert jnp.array_equal(out, jnp.array([0.0, 2.0, 3.0]))
print(out)

def add_scalar(vector, scalar):
    return vector + scalar

vectors = jnp.arange(12).reshape(4, 3)     # 4 个向量，每个长度为 3
scalar_val = 100.0

# 映射 vectors 的第 0 轴，广播 scalar（不映射）
vectorized_add_scalar = jax.vmap(add_scalar, in_axes=(0, None))
result = vectorized_add_scalar(vectors, scalar_val)
print("in_axes=(0, None):", result.shape)
print(result)

# 每个向量加上各自的标量：两个参数都沿第 0 轴映射
scalars = jnp.array([100., 200., 300., 400.])
result_batched_scalars = jax.vmap(add_scalar, in_axes=(0, 0))(vectors, scalars)
print("in_axes=(0, 0):", result_batched_scalars.shape)
print(result_batched_scalars)

# 沿非 0 轴映射：数据转置后，每一列是一个向量
vectors_T = vectors.T                       # 形状 (3, 4)
result_axis1 = jax.vmap(add_scalar, in_axes=(1, None))(vectors_T, scalar_val)
print("in_axes=(1, None):", result_axis1.shape, "（输出的批次维度仍在第 0 轴）")
print(result_axis1)
assert result_axis1.shape == (4, 3) and jnp.array_equal(result_axis1, result)

def process_vector(v):
    return v * 2

# 默认 out_axes=0
out0 = jax.vmap(process_vector, in_axes=0, out_axes=0)(vectors)
print("out_axes=0:", out0.shape)
print(out0)

# out_axes=1：原向量维度（大小 3）成为输出第 0 轴，批次维度（大小 4）放到第 1 轴
out1 = jax.vmap(process_vector, in_axes=0, out_axes=1)(vectors)
print("out_axes=1:", out1.shape)
print(out1)

# 多个返回值（pytree）：为不同输出指定不同的 out_axes
def process_vector_pytree(v):
    return {'sum': v.sum(), 'doubled': v * 2}

res = jax.vmap(process_vector_pytree, in_axes=0, out_axes={'sum': 0, 'doubled': 1})(vectors)
print("sum:", res['sum'].shape, res['sum'])
print("doubled:", res['doubled'].shape)
print(res['doubled'])

k1, k2 = jax.random.split(jax.random.key(1))
W = jax.random.normal(k1, (3, 4))     # 共享权重
X = jax.random.normal(k2, (5, 4))     # 5 个样本

def linear(W, x):                      # 单样本：矩阵乘向量
    return W @ x

Y = jax.vmap(linear, in_axes=(None, 0))(W, X)
print("Y 形状:", Y.shape, "| 与手写批量版本一致:", jnp.allclose(Y, X @ W.T, atol=1e-5))
print("in_axes=-1（按列存放的数据）:", jnp.allclose(jax.vmap(linear, in_axes=(None, -1))(W, X.T), Y))

params = {"w": jnp.array([1.0, 2.0, 3.0]),          # 共享权重：形状 (3,)
          "b": jnp.array([0.0, 10.0, 20.0, 30.0])}  # 每个样本各自的偏置：形状 (4,)
xb = jnp.arange(12.0).reshape(4, 3)                 # 4 个样本

def predict_d(p, x):
    return jnp.dot(p["w"], x) + p["b"]

# 对 params 用 dict 形式的 in_axes：w 不映射，b 沿第 0 轴映射
out = jax.vmap(predict_d, in_axes=({"w": None, "b": 0}, 0))(params, xb)
print(out)
print("手动核对:", xb @ params["w"] + params["b"])

# 结构不匹配（键名写错）会报错
try:
    jax.vmap(predict_d, in_axes=({"w": None, "bias": 0}, 0))(params, xb)
except Exception as e:
    print(type(e).__name__, ":", str(e)[:420])

add = lambda a, b: a + b

# 1) 关键字参数：元组 in_axes 不对应关键字参数，会报错并给出提示
try:
    jax.vmap(add, in_axes=(0, None))(jnp.ones(3), b=jnp.ones(3))
    raised = False
except ValueError as e:
    raised = "keyword arguments" in str(e)
    print("关键字参数报错:", str(e)[:260], "...")
assert raised, "元组 in_axes 不对应关键字参数，应当报错"
# 关键字参数会被沿第 0 轴映射：传入长度 3 的 b，得到逐元素相加，而不是“整体广播”
assert jnp.array_equal(jax.vmap(add)(jnp.arange(3.0), b=jnp.arange(3.0)), 2 * jnp.arange(3.0))

# 2) out_axes=None 但输出依赖批数据 -> 报错
try:
    jax.vmap(lambda x: 2 * x, out_axes=None)(jnp.arange(3.0))
except ValueError as e:
    print("out_axes=None 报错:", e)
# 输出不依赖批数据时，out_axes=None 合法
print("out_axes=None 合法用法:", jax.vmap(lambda x, c: (x * c, c), in_axes=(0, None), out_axes=(0, None))(jnp.arange(3.0), 5.0))

# 3) axis_size：没有任何被映射的输入时必须显式给出
try:
    jax.vmap(lambda c: c * jnp.ones(2), in_axes=None)(3.0)
except ValueError as e:
    print("缺少 axis_size 报错:", e)
print("axis_size=4:\n", jax.vmap(lambda c: c * jnp.ones(2), in_axes=None, axis_size=4)(3.0))

# 4) sum_match=True：out_axes=None 的输出沿批维度求和
sm = jax.vmap(lambda x: 2.0 * x, out_axes=None, sum_match=True)(jnp.arange(3.0))
assert sm == 6.0                                  # 2*(0+1+2)
print("sum_match:", sm)

# 场景 1：映射部分参数，广播其他参数
def scale_vector(vector, scalar):
    return vector * scalar

vector = jnp.arange(3.)
scalar = 2.0
print("单次缩放:", scale_vector(vector, scalar))          # [0. 2. 4.]

batch_of_vectors = jnp.arange(9.).reshape((3, 3))
single_scalar = 2.0
vectorized_scale = jax.vmap(scale_vector, in_axes=(0, None))
scaled = vectorized_scale(batch_of_vectors, single_scalar)
print("批量缩放:\n", scaled)
assert jnp.array_equal(scaled, batch_of_vectors * 2.0)

# ❌ 错误写法：以为默认只映射第一个参数，于是省略了 in_axes
try:
    jax.vmap(scale_vector)(batch_of_vectors, 2.0)
    raised = False
except ValueError as e:
    raised = True
    print("实测 ValueError:", e)
assert raised, "默认 in_axes=0 也会去映射标量参数，应当报错"

# ✅ 推荐写法：默认 in_axes=0 映射所有位置参数，共享的标量必须显式写 None
ok = jax.vmap(scale_vector, in_axes=(0, None))(batch_of_vectors, 2.0)
assert jnp.array_equal(ok, batch_of_vectors * 2.0)

# 验证“默认映射所有位置参数”：两个参数都带批维度时，默认 in_axes 会同时映射两者
def add_vectors(vec1, vec2):
    return vec1 + vec2
b1 = jnp.arange(9.).reshape(3, 3)
b2 = jnp.arange(3.)[:, None] * jnp.ones((3, 3))    # 第 i 行全是 i
default_out = jax.vmap(add_vectors)(b1, b2)
assert jnp.array_equal(default_out, b1 + b2)       # 第 i 个样本 = b1[i] + b2[i]，第二个参数也被映射了
print("✔ 断言通过：默认 in_axes=0 映射的是所有位置参数")

# 场景 2：多个参数沿相同轴映射
batch1 = jnp.arange(9.).reshape((3, 3))
batch2 = jnp.ones((3, 3)) * 10
vectorized_add = jax.vmap(add_vectors, in_axes=(0, 0))
added = vectorized_add(batch1, batch2)
print(added)
assert jnp.array_equal(added, batch1 + 10)

# 场景 3（前半）：行 × 滤波向量，映射矩阵的行，广播滤波向量
def apply_filter(row_vector, filter_vector):
    return row_vector * filter_vector

matrix = jnp.arange(12.).reshape((3, 4))
filter_v = jnp.array([1., 10., 100., 1000.])
filtered = jax.vmap(apply_filter, in_axes=(0, None))(matrix, filter_v)
print("行滤波 in_axes=(0, None):\n", filtered)
assert filtered.shape == (3, 4)

# ❌ 容易误解的写法：以为结果是 3x4、“映射维度放回第 1 轴”
def scale_column(column_vector, scalar):
    return column_vector * scalar

scalars = jnp.array([1., 10., 100., 1000.])
if matrix.shape[1] == scalars.shape[0]:
    scaled_cols = jax.vmap(scale_column, in_axes=(1, 0))(matrix, scalars)
    print("实测形状:", scaled_cols.shape)
    print(scaled_cols)

# 验证：默认 out_axes=0，批次维度（4 列）在输出第 0 轴 -> (4, 3)，是期望结果的转置
assert scaled_cols.shape == (4, 3)
assert jnp.allclose(scaled_cols, (matrix * scalars).T)

# ✅ 推荐写法：用 out_axes=1 把批次维度放回第 1 轴，才能得到想要的 3x4 布局
scaled_cols_fixed = jax.vmap(scale_column, in_axes=(1, 0), out_axes=1)(matrix, scalars)
print("修正后形状:", scaled_cols_fixed.shape)
print(scaled_cols_fixed)
assert scaled_cols_fixed.shape == (3, 4)
assert jnp.allclose(scaled_cols_fixed, matrix * scalars)    # 与广播写法一致
print("✔ 断言通过")

f2 = lambda a, b: a + b

cases = {
    "批大小不一致 (3 vs 4)": lambda: jax.vmap(f2)(jnp.ones(3), jnp.ones(4)),
    "长度 1 不会被自动广播": lambda: jax.vmap(f2)(jnp.ones((3, 2)), jnp.ones((1, 2))),
    "in_axes 长度与参数个数不符": lambda: jax.vmap(f2, in_axes=(0,))(jnp.ones(3), jnp.ones(3)),
    "映射轴超出数组维数": lambda: jax.vmap(lambda a: a, in_axes=1)(jnp.ones(3)),
}
n_raised = 0
for name, fn in cases.items():
    try:
        fn()
        print(f"[{name}] 没有报错?!")
    except ValueError as e:
        n_raised += 1
        print(f"[{name}]\n  ValueError: {e}\n")
assert n_raised == len(cases)

# 单个矩阵乘单个向量
def matvec_multiply(matrix, vector):
    return jnp.dot(matrix, vector)

matrix = jnp.arange(6.).reshape(2, 3)
vector = jnp.arange(3.)
print("单次 matvec:", matvec_multiply(matrix, vector))        # [ 5. 14.]，形状 (2,)

# 单层 vmap：矩阵批与向量批按位置配对
batch_of_matrices = jnp.arange(18.).reshape(3, 2, 3)          # 3 个 2x3 矩阵
batch_of_vectors = jnp.arange(9.).reshape(3, 3)               # 3 个 3 维向量
batched_matvec = jax.vmap(matvec_multiply, in_axes=(0, 0))
print("配对 matvec 形状:", batched_matvec(batch_of_matrices, batch_of_vectors).shape)   # (3, 2)

# 嵌套 vmap：每个矩阵 × 每个向量（所有组合）
batch_of_vectors_outer = jnp.arange(12.).reshape(4, 3)        # 4 个向量

# 内层 in_axes=(None, 0)：矩阵固定，遍历所有向量
# 外层 in_axes=(0, None)：遍历所有矩阵，向量批整体共享
pairwise_matvec = jax.vmap(
    jax.vmap(matvec_multiply, in_axes=(None, 0)),
    in_axes=(0, None),
)
result = pairwise_matvec(batch_of_matrices, batch_of_vectors_outer)
print("两两 matvec 形状:", result.shape)                       # (3, 4, 2)
assert result.shape == (3, 4, 2)

# 核对：result[i, j] 应该等于第 i 个矩阵乘第 j 个向量
assert jnp.allclose(result[1, 2], matvec_multiply(batch_of_matrices[1], batch_of_vectors_outer[2]))
print("✔ result[1, 2] == 第 1 个矩阵 @ 第 2 个向量")

# 两两距离（欧氏距离的平方）
def euclidean_distance_sq(point_a, point_b):
    return jnp.sum((point_a - point_b)**2)

points_a = jnp.array([[1., 0.], [0., 1.], [-1., 0.]])                 # (3, 2)
points_b = jnp.array([[2., 2.], [-2., 2.], [2., -2.], [-2., -2.]])    # (4, 2)

inner_map = jax.vmap(euclidean_distance_sq, in_axes=(None, 0))      # 一个 a 对所有 b
pairwise_distance_sq = jax.vmap(inner_map, in_axes=(0, None))       # 所有 a
# 等价的一行写法
pairwise_distance_sq_direct = jax.vmap(
    jax.vmap(euclidean_distance_sq, in_axes=(None, 0)), in_axes=(0, None))

distance_matrix = pairwise_distance_sq(points_a, points_b)
print("距离矩阵形状:", distance_matrix.shape)                       # (3, 4)
print(distance_matrix)
assert jnp.allclose(distance_matrix, pairwise_distance_sq_direct(points_a, points_b))
assert jnp.array_equal(distance_matrix, jnp.array([[5., 13., 5., 13.], [5., 5., 13., 13.], [13., 5., 13., 5.]]))
print("手工核对 [0, 0]:", euclidean_distance_sq(points_a[0], points_b[0]))   # (1-2)^2 + (0-2)^2 = 5

# 嵌套 vmap 与 jit、grad 组合
jitted_pairwise_distance_sq = jax.jit(pairwise_distance_sq_direct)
print("jit 版本一致:", jnp.allclose(jitted_pairwise_distance_sq(points_a, points_b), distance_matrix))

def sum_of_distances(p_a, p_b):
    dist_matrix = pairwise_distance_sq_direct(p_a, p_b)
    return jnp.sum(dist_matrix)

grad_sum_dist_wrt_a = jax.grad(sum_of_distances, argnums=0)
gradients_a = grad_sum_dist_wrt_a(points_a, points_b)
print("对 points_a 的梯度形状:", gradients_a.shape)                  # (3, 2)，与 points_a 相同
print(gradients_a)
# 解析解核对：d/da_i sum_j |a_i - b_j|^2 = sum_j 2 (a_i - b_j)
assert jnp.allclose(gradients_a, 2 * (points_b.shape[0] * points_a - points_b.sum(0)))
print("✔ 梯度与解析解一致")

k1, k2 = jax.random.split(jax.random.key(3))
Xs = jax.random.normal(k1, (4, 3))    # M=4
Ys = jax.random.normal(k2, (5, 3))    # N=5

G = jax.vmap(jax.vmap(jnp.dot, in_axes=(None, 0)), in_axes=(0, None))(Xs, Ys)
G2 = jax.vmap(jax.vmap(jnp.dot, in_axes=(0, None)), in_axes=(None, 0))(Xs, Ys)
print("G:", G.shape, "| 交换顺序后 G2:", G2.shape, "| G2 == G.T:", jnp.allclose(G2, G.T))
print("等于 X @ Y.T:", jnp.allclose(G, Xs @ Ys.T, atol=1e-5))

data = jnp.ones((2, 3, 4))            # (B, T, d)
print("双重批维度，逐位置求范数:", jax.vmap(jax.vmap(jnp.linalg.norm))(data).shape)

def dist(x, y):
    return jnp.sqrt(jnp.sum((x - y) ** 2))

pdist = jax.vmap(jax.vmap(dist, in_axes=(None, 0)), in_axes=(0, None))
D = pdist(Xs, Ys)

# 对照：广播写法（展开公式），需要 clip 防止出现负数
sq = jnp.sum(Xs**2, 1)[:, None] + jnp.sum(Ys**2, 1)[None, :] - 2 * Xs @ Ys.T
D_ref = jnp.sqrt(jnp.clip(sq, 0.0))
assert jnp.allclose(D, D_ref, atol=1e-4)
print("距离矩阵形状:", D.shape, "| 与展开公式一致")

# 梯度陷阱：x == y 时梯度为 nan（修法“双 where”见第 3 章 3.8）
g_naive = jax.grad(dist)(jnp.ones(3), jnp.ones(3))
print("朴素 dist 在 x==y 处的梯度:", g_naive)
assert jnp.all(jnp.isnan(g_naive))

k1, k2 = jax.random.split(jax.random.key(4))
Ab = jax.random.normal(k1, (64, 128, 96))
Bb = jax.random.normal(k2, (64, 96, 80))

bmm_vmap = jax.jit(jax.vmap(jnp.matmul))
bmm_einsum = jax.jit(lambda a, b: jnp.einsum("bnk,bkm->bnm", a, b))
bmm_at = jax.jit(lambda a, b: a @ b)

r1, r2, r3 = bmm_vmap(Ab, Bb), bmm_einsum(Ab, Bb), bmm_at(Ab, Bb)
assert jnp.allclose(r1, r2, atol=1e-3) and jnp.allclose(r1, r3, atol=1e-3)
print("结果形状:", r1.shape, "| 三者一致")
for name, fn in [("vmap(matmul)", bmm_vmap), ("einsum", bmm_einsum), ("a @ b", bmm_at)]:
    print(f"{name:14s} {bench(fn, Ab, Bb):.3f} ms")

# 看 jaxpr：都是带批维度的 dot_general
jp_bmm = [jaxpr_str(jax.make_jaxpr(f)(Ab[:2], Bb[:2])) for f in
          (jax.vmap(jnp.matmul), lambda a, b: jnp.einsum("bnk,bkm->bnm", a, b), lambda a, b: a @ b)]
assert jp_bmm[0] == jp_bmm[1] == jp_bmm[2]    # 验证：三种写法的 jaxpr 逐字相同
print(jax.make_jaxpr(jax.vmap(jnp.matmul))(Ab[:2], Bb[:2]))
print(jax.make_jaxpr(lambda a, b: jnp.einsum("bnk,bkm->bnm", a, b))(Ab[:2], Bb[:2]))

import time

# 简单线性模型
def predict(params, x):
    w, b = params
    return w * x + b

params = (jnp.array(2.0), jnp.array(1.0))     # (w, b)
batch_x = jnp.arange(10000.0)

# 在 batch_x 的第 0 轴上映射，params 保持不变
batched_predict_vmap = jax.vmap(predict, in_axes=(None, 0))
jitted_batched_predict = jax.jit(batched_predict_vmap)

t0 = time.perf_counter()
result_vmap = batched_predict_vmap(params, batch_x).block_until_ready()
print(f"仅 vmap 的执行时间: {(time.perf_counter() - t0) * 1e3:.3f} ms")

t0 = time.perf_counter()
result_jit_vmap = jitted_batched_predict(params, batch_x).block_until_ready()
print(f"jit(vmap) 执行时间（首次运行，含编译）: {(time.perf_counter() - t0) * 1e3:.3f} ms")

t0 = time.perf_counter()
result_jit_vmap = jitted_batched_predict(params, batch_x).block_until_ready()
print(f"jit(vmap) 执行时间（第二次运行）: {(time.perf_counter() - t0) * 1e3:.3f} ms")

assert jnp.allclose(result_vmap, result_jit_vmap)
print("结果匹配")

# 上面的单次计时只为展示“首次运行含编译”；比较稳定耗时要预热后多次取平均
print(f"预热后平均: vmap {bench(batched_predict_vmap, params, batch_x):.3f} ms | "
      f"jit(vmap) {bench(jitted_batched_predict, params, batch_x):.3f} ms")

# grad 与 vmap：逐样本梯度
def squared_error(params, x, y_true):
    w, b = params
    y_pred = w * x + b
    loss = (y_pred - y_true)**2
    return loss

# 单个 (x, y_true) 的梯度函数：返回 (dL/dw, dL/db)
grad_loss_single = jax.grad(squared_error, argnums=0)

params = (2.0, 1.0)
batch_x = jnp.linspace(0, 1, 5)
batch_y_true = 2.5 * batch_x + 0.5                # 真实关系：w=2.5, b=0.5

grad_loss_batch = jax.vmap(grad_loss_single, in_axes=(None, 0, 0))
batch_gradients = grad_loss_batch(params, batch_x, batch_y_true)

print("params:", params)
print("batch_x:", batch_x)
print("batch_y_true:", batch_y_true)
print("逐样本 dw:", batch_gradients[0], batch_gradients[0].shape)   # (5,)
print("逐样本 db:", batch_gradients[1], batch_gradients[1].shape)   # (5,)
# 解析解核对：dL/dw = 2 (w x + b - y) x，dL/db = 2 (w x + b - y)
resid = 2.0 * batch_x + 1.0 - batch_y_true
assert jnp.allclose(batch_gradients[0], 2 * resid * batch_x) and jnp.allclose(batch_gradients[1], 2 * resid)

# 批次平均损失 -> 平均梯度
def mean_squared_error(params, batch_x, batch_y_true):
    batch_predict = jax.vmap(predict, in_axes=(None, 0))
    predictions = batch_predict(params, batch_x)
    return jnp.mean((predictions - batch_y_true)**2)

grad_mean_loss = jax.grad(mean_squared_error, argnums=0)
mean_batch_gradient = grad_mean_loss(params, batch_x, batch_y_true)
print("平均梯度:", mean_batch_gradient)
print("形状:", jax.tree.map(jnp.shape, mean_batch_gradient))        # 每个分量都是标量 ()

# 平均梯度 == 逐样本梯度的平均
assert jnp.allclose(mean_batch_gradient[0], jnp.mean(batch_gradients[0]))
assert jnp.allclose(mean_batch_gradient[1], jnp.mean(batch_gradients[1]))
print("✔ 平均损失的梯度 == 逐样本梯度的平均")

# ❌ 旧写法（过时）：PRNGKey 是旧式原始 key；且同一个 key 被用于两次采样，x 与“噪声”高度相关
key = jax.random.PRNGKey(42)
print("PRNGKey 实测:", key.dtype, key.shape, "（原始 uint32 数组，不是带类型的 key）")
large_batch_size = 100000
large_batch_x_bad = jax.random.uniform(key, (large_batch_size,))
noise_bad = jax.random.normal(key, (large_batch_size,))          # 又用了同一个 key
rank = lambda a: np.argsort(np.argsort(np.asarray(a)))
corr = np.corrcoef(np.asarray(large_batch_x_bad), np.asarray(noise_bad))[0, 1]
rank_corr = np.corrcoef(rank(large_batch_x_bad), rank(noise_bad))[0, 1]
print(f"实测：x 与噪声的线性相关系数 {corr:.4f}，秩相关系数 {rank_corr:.4f}")
assert corr > 0.95 and rank_corr > 0.999          # 验证：同一个 key 生成的 uniform 与 normal 高度相关

# ✅ 推荐写法：jax.random.key 创建带类型的 key，先 split 再分别使用
key = jax.random.key(42)
key_x, key_noise = jax.random.split(key)
print("jax.random.key 实测:", key.dtype, key.shape)
large_batch_x = jax.random.uniform(key_x, (large_batch_size,))
noise = jax.random.normal(key_noise, (large_batch_size,))
corr_ok = np.corrcoef(np.asarray(large_batch_x), np.asarray(noise))[0, 1]
print(f"split 之后的相关系数: {corr_ok:.4f}")
assert abs(corr_ok) < 0.02                        # 独立：相关系数接近 0
large_batch_y_true = 2.5 * large_batch_x + 0.5 + 0.1 * noise

# 课程代码（续）：jit(vmap(grad(f)))，大批量的逐样本梯度（数据来自上面的推荐写法）
vmap_grad_loss = jax.vmap(grad_loss_single, in_axes=(None, 0, 0))
jit_vmap_grad_loss = jax.jit(vmap_grad_loss)
params = (2.0, 1.0)

t0 = time.perf_counter()
per_example_grads = jit_vmap_grad_loss(params, large_batch_x, large_batch_y_true)
jax.block_until_ready(per_example_grads)
print(f"首次运行（含编译）: {(time.perf_counter() - t0) * 1e3:.2f} ms")

t0 = time.perf_counter()
per_example_grads = jit_vmap_grad_loss(params, large_batch_x, large_batch_y_true)
jax.block_until_ready(per_example_grads)
print(f"第二次运行: {(time.perf_counter() - t0) * 1e3:.2f} ms")
print("dw 形状:", per_example_grads[0].shape, "| db 形状:", per_example_grads[1].shape)
assert per_example_grads[0].shape == per_example_grads[1].shape == (large_batch_size,)

f_feat = lambda x: jnp.sin(x) @ jnp.cos(x)
jp = jaxpr_str(jax.make_jaxpr(jax.vmap(jax.jit(f_feat)))(jnp.ones((4, 3))))
print(jp)
# 验证：jit 内部已经是整批 (4, 3) 的运算
inner = jp[jp.index("jaxpr={"):]
assert "f32[4,3] = sin" in inner and "f32[4] = dot_general" in inner

for shape in [(512, 256), (4096, 1024)]:
    xb = jax.random.normal(jax.random.key(6), shape)
    a_ = jax.jit(jax.vmap(f_feat))
    b_ = jax.vmap(jax.jit(f_feat))
    assert jnp.allclose(a_(xb), b_(xb), rtol=1e-4, atol=1e-3)
    print(f"批 {shape}: jit(vmap(f)) {bench(a_, xb):.3f} ms | vmap(jit(f)) {bench(b_, xb):.3f} ms")

# 准备：两层 MLP + 回归数据，并故意把 5 个样本的标签改坏（后面用来找异常样本）
def init_mlp_params(key, d_in, d_hidden):
    k1, k2 = jax.random.split(key)
    return {"W1": jax.random.normal(k1, (d_hidden, d_in)) / jnp.sqrt(d_in), "b1": jnp.zeros(d_hidden),
            "w2": jax.random.normal(k2, (d_hidden,)) / jnp.sqrt(d_hidden), "b2": jnp.array(0.0)}

def mlp_loss(p, x, y):                               # 单样本损失（标量）
    h = jnp.tanh(p["W1"] @ x + p["b1"])
    return (jnp.dot(p["w2"], h) + p["b2"] - y) ** 2

B, D_IN, H = 256, 8, 32
kp, kx, kn = jax.random.split(jax.random.key(100), 3)
pe_params = init_mlp_params(kp, D_IN, H)
X_pe = jax.random.normal(kx, (B, D_IN))
Y_pe = jnp.sin(X_pe[:, 0]) + 0.1 * jax.random.normal(kn, (B,))
bad_idx = jnp.array([3, 50, 99, 150, 222])           # 被“标错”的样本
Y_pe = Y_pe.at[bad_idx].add(20.0)

P = sum(l.size for l in jax.tree.leaves(pe_params))
print("参数量 P =", P, "| 批大小 B =", B)

# 1) 对拍：vmap(grad) vs Python 循环逐个求梯度，并与 jacrev（雅可比视角）比较
per_ex_grad = jax.jit(jax.vmap(jax.grad(mlp_loss), in_axes=(None, 0, 0)))
single_grad = jax.jit(jax.grad(mlp_loss))

def loop_per_ex_grad(p, X, Y):
    gs = [single_grad(p, X[i], Y[i]) for i in range(X.shape[0])]
    return jax.tree.map(lambda *t: jnp.stack(t), *gs)

g_vmap = per_ex_grad(pe_params, X_pe, Y_pe)
g_loop = loop_per_ex_grad(pe_params, X_pe, Y_pe)
g_jac = jax.jacrev(lambda p: jax.vmap(mlp_loss, in_axes=(None, 0, 0))(p, X_pe, Y_pe))(pe_params)

tree_allclose = lambda a, b, tol=1e-5: all(jax.tree.leaves(jax.tree.map(lambda u, v: bool(jnp.allclose(u, v, atol=tol)), a, b)))
assert tree_allclose(g_vmap, g_loop)
assert tree_allclose(g_vmap, g_jac)
print("逐样本梯度形状:", jax.tree.map(jnp.shape, g_vmap))
print("✔ vmap(grad) == 循环逐个求梯度 == jacrev")

t_loop = bench(lambda p, X, Y: loop_per_ex_grad(p, X, Y)["W1"], pe_params, X_pe, Y_pe, n=2)
t_vmap = bench(lambda p, X, Y: per_ex_grad(p, X, Y)["W1"], pe_params, X_pe, Y_pe)
print(f"Python 循环: {t_loop:.2f} ms | jit(vmap(grad)): {t_vmap:.3f} ms | 加速约 {t_loop / t_vmap:.0f} 倍")

# 2) 逐样本梯度的均值 == 平均损失的梯度
def mean_mlp_loss(p, X, Y):
    return jnp.mean(jax.vmap(mlp_loss, in_axes=(None, 0, 0))(p, X, Y))

g_mean_loss = jax.jit(jax.grad(mean_mlp_loss))(pe_params, X_pe, Y_pe)
g_mean_of_per_ex = jax.tree.map(lambda t: t.mean(axis=0), g_vmap)
assert tree_allclose(g_mean_loss, g_mean_of_per_ex)
print("✔ mean_i g_i == grad(mean loss)")

# 对照：grad(vmap(loss)) 不能直接用——输出不是标量
try:
    jax.grad(jax.vmap(mlp_loss, in_axes=(None, 0, 0)))(pe_params, X_pe, Y_pe)
    raised = False
except TypeError as e:
    raised = True
    print("grad(vmap(loss)) -> TypeError:", str(e)[:90])
assert raised

# 3) 应用 A：DP-SGD 式裁剪 —— 每个样本的梯度先按全局 L2 范数裁剪到 C，再平均并加高斯噪声
def per_example_norms(g):
    # 每个叶子展平到 (B, -1) 后求平方和，再把各叶子相加
    return jnp.sqrt(sum(jnp.sum(l.reshape(l.shape[0], -1) ** 2, axis=1) for l in jax.tree.leaves(g)))

def clip_per_example(g, C):
    scale = jnp.minimum(1.0, C / (per_example_norms(g) + 1e-12))
    return jax.tree.map(lambda l: l * scale.reshape((-1,) + (1,) * (l.ndim - 1)), g)

C, sigma = 1.0, 0.5
g_clipped = clip_per_example(g_vmap, C)
norms_before, norms_after = per_example_norms(g_vmap), per_example_norms(g_clipped)
assert jnp.all(norms_after <= C + 1e-5)                            # 裁剪后每个样本的范数都不超过 C
assert jnp.allclose(norms_after, jnp.minimum(norms_before, C), atol=1e-4)

leaves, treedef = jax.tree.flatten(g_clipped)
noise_keys = jax.random.split(jax.random.key(7), len(leaves))      # 每个叶子用独立的 key
dp_grad = jax.tree.unflatten(treedef, [l.mean(0) + sigma * C / B * jax.random.normal(k, l.shape[1:])
                                       for l, k in zip(leaves, noise_keys)])
print(f"裁剪前最大范数 {norms_before.max():.2f} -> 裁剪后最大范数 {norms_after.max():.2f}")
print("DP 梯度的形状与参数相同:", jax.tree.map(jnp.shape, dp_grad) == jax.tree.map(jnp.shape, pe_params))

# 4) 应用 B：用逐样本梯度范数找出异常样本（这里是被“标错”的 5 个样本）
norms = per_example_norms(g_vmap)
top5 = jnp.argsort(norms)[-5:]
print("梯度范数最大的 5 个样本:", jnp.sort(top5), "| 真正标错的样本:", bad_idx)
print("异常样本范数中位数 %.2f vs 正常样本范数中位数 %.2f" % (
    jnp.median(norms[bad_idx]), jnp.median(jnp.delete(norms, bad_idx))))
assert set(np.asarray(top5).tolist()) == set(np.asarray(bad_idx).tolist())

# 5) 应用 C：per-example loss 加权 —— 加权损失的梯度 == 逐样本梯度的加权和
weights = jnp.ones(B).at[bad_idx].set(0.0)            # 例：把找到的异常样本权重设为 0
weights = weights / weights.sum()

def weighted_loss(p, X, Y, w):
    return jnp.sum(w * jax.vmap(mlp_loss, in_axes=(None, 0, 0))(p, X, Y))

g_weighted = jax.jit(jax.grad(weighted_loss))(pe_params, X_pe, Y_pe, weights)
g_weighted_from_per_ex = jax.tree.map(lambda l: jnp.tensordot(weights, l, axes=1), g_vmap)
assert tree_allclose(g_weighted, g_weighted_from_per_ex)
print("✔ grad(sum_i w_i * loss_i) == sum_i w_i * g_i")
print("去掉异常样本前后，平均梯度范数: %.3f -> %.3f" % (
    per_example_norms(jax.tree.map(lambda l: l[None], g_mean_loss))[0],
    per_example_norms(jax.tree.map(lambda l: l[None], g_weighted))[0]))

# 6) 内存：逐样本梯度大小 = B × P；对比编译后的输出与临时内存
n_elems = sum(l.size for l in jax.tree.leaves(g_vmap))
assert n_elems == B * P
print(f"逐样本梯度: {n_elems} 个元素 = B({B}) × P({P}) ≈ {n_elems * 4 / 2**20:.2f} MiB；平均梯度只有 {P * 4 / 2**10:.2f} KiB")

for name, fn in [("jit(vmap(grad))", per_ex_grad), ("jit(grad(mean))", jax.jit(jax.grad(mean_mlp_loss)))]:
    ma = fn.lower(pe_params, X_pe, Y_pe).compile().memory_analysis()
    print(f"{name:16s} 输出 {ma.output_size_in_bytes / 2**10:8.1f} KiB | 临时 {ma.temp_size_in_bytes / 2**10:8.1f} KiB")

# 只要范数时：lax.map 分块计算，每块只返回范数，峰值内存与块大小成正比
@jax.jit
def norms_chunked(p, X, Y):
    one = lambda xy: per_example_norms(jax.tree.map(lambda l: l[None], jax.grad(mlp_loss)(p, xy[0], xy[1])))[0]
    return lax.map(one, (X, Y), batch_size=32)

nc = norms_chunked(pe_params, X_pe, Y_pe)
assert jnp.allclose(nc, norms, rtol=1e-4, atol=1e-5)
ma = norms_chunked.lower(pe_params, X_pe, Y_pe).compile().memory_analysis()
print(f"lax.map 分块只算范数: 输出 {ma.output_size_in_bytes / 2**10:.1f} KiB | 临时 {ma.temp_size_in_bytes / 2**10:.1f} KiB，结果一致 ✔")