第 10 章　随机数与分片

# ⚠️ 必须在 import jax 之前设置：在 CPU 上模拟 8 个设备（若已导入过 jax，请重启内核）
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import sys, re, io, json, time, socket, subprocess, warnings
import random as py_random            # Python 标准库的 random，起别名以免和 jax.random 混淆
from collections import Counter

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.sharding import NamedSharding, PartitionSpec as P
from scipy import stats
import flax
from flax import nnx

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| Flax", flax.__version__)
print("设备  ", jax.devices())
assert jax.__version__ == "0.12.0.1"
assert jax.device_count() == 8


def bench(f, *args, warmup=2, runs=10):
    """计时：先预热 warmup 次（含编译），再运行 runs 次取中位数，返回毫秒。只用于打印。"""
    for _ in range(warmup):
        jax.block_until_ready(f(*args))
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args))
        ts.append(time.perf_counter() - t0)
    return float(np.median(ts)) * 1e3


COLLECTIVE_OPS = ["all-reduce", "all-gather", "reduce-scatter", "all-to-all", "collective-permute"]

def collectives(hlo_text):
    """统计编译后 HLO 文本中各类集合通信出现的次数（含 *-start 异步形式）。"""
    return {op: len(re.findall(op + r"(?:-start)?\(", hlo_text)) for op in COLLECTIVE_OPS}


def run_py(code, *args, env=None, timeout=600):
    """在全新的 Python 子进程中执行 code，返回它最后一行 'JSON:' 输出解析出的对象。"""
    full_env = dict(os.environ, **(env or {}))
    r = subprocess.run([sys.executable, "-c", code, *map(str, args)],
                       capture_output=True, text=True, env=full_env, timeout=timeout)
    assert r.returncode == 0, r.stderr[-3000:]
    lines = [l for l in r.stdout.splitlines() if l.startswith("JSON:")]
    assert lines, r.stdout[-2000:] + r.stderr[-2000:]
    return json.loads(lines[-1][5:])


def err_line(e):
    """异常的类型 + 第一行信息（演示错误写法时用）。"""
    msg = str(e).strip().splitlines()[0] if str(e).strip() else ""
    return f"{type(e).__name__}: {msg[:170]}"

print("公共工具已定义")

# NumPy 的全局状态：同一个种子，只是交换了两个函数的调用顺序
def init_weights_np():
    return np.random.normal(size=3)

def sample_noise_np():
    return np.random.normal(size=3)

np.random.seed(0); w_a = init_weights_np(); n_a = sample_noise_np()
np.random.seed(0); n_b = sample_noise_np(); w_b = init_weights_np()   # 顺序反过来
print("先 init 再 noise 的权重:", w_a)
print("先 noise 再 init 的权重:", w_b)

assert not np.array_equal(w_a, w_b)       # 权重变了
assert np.array_equal(w_a, n_b)            # 而且“权重”和“噪声”互换了身份
print("✓ NumPy 全局状态：结果依赖调用顺序")

# JAX：随机数只由 key 决定，调用顺序无关
def init_weights(key):
    return jax.random.normal(key, (3,))

def sample_noise(key):
    return jax.random.normal(key, (3,))

k_w, k_n = jax.random.split(jax.random.key(0))
w1 = init_weights(k_w); n1 = sample_noise(k_n)
n2 = sample_noise(k_n); w2 = init_weights(k_w)      # 顺序反过来
assert jnp.array_equal(w1, w2) and jnp.array_equal(n1, n2)

# key 是普通参数：jit 只追踪 / 编译一次，换 key 不会重新编译
trace_count = 0
def noisy(key, x):
    global trace_count
    trace_count += 1                  # 只在追踪时执行
    return x + jax.random.normal(key, x.shape)

noisy_jit = jax.jit(noisy)
x0 = jnp.zeros(3)
outs = [noisy_jit(k, x0) for k in jax.random.split(jax.random.key(1), 4)]
print("4 个不同 key 的结果互不相同:", len({tuple(np.asarray(o).tolist()) for o in outs}) == 4,
      "| 追踪次数 =", trace_count)
assert trace_count == 1
assert len({tuple(np.asarray(o).tolist()) for o in outs}) == 4
print("✓ JAX：顺序无关、可复现，key 作为参数不触发重新编译")

k = jax.random.key(0)
raw = jax.random.PRNGKey(0)

# 判断“是不是类型化 key”的标准写法：jax.dtypes.prng_key 是所有 key dtype 的抽象父类
print("key(0).dtype =", k.dtype, "| 是 prng_key:", jnp.issubdtype(k.dtype, jax.dtypes.prng_key))
print("PRNGKey(0).dtype =", raw.dtype, "| 是 prng_key:", jnp.issubdtype(raw.dtype, jax.dtypes.prng_key))
print("key 记得自己的算法:", jax.random.key_impl(k))
assert jnp.issubdtype(k.dtype, jax.dtypes.prng_key)
assert not jnp.issubdtype(raw.dtype, jax.dtypes.prng_key)
assert str(jax.random.key_impl(k)) == "threefry2x32"

# 形状语义：key.shape 是 key 的个数，key_data 多一个尾部维度
print(f"\n{'表达式':28s} {'key.shape':>10s} {'key_data.shape':>15s}")
for name, kk in [("key(0)", k),
                 ("split(key, 4)", jax.random.split(k, 4)),
                 ("split(key, (2, 3))", jax.random.split(k, (2, 3))),
                 ("key(0, impl='rbg')", jax.random.key(0, impl="rbg")),
                 ("split(rbg_key, (2, 3))", jax.random.split(jax.random.key(0, impl="rbg"), (2, 3))),
                 ("key(0, impl='philox2x32')", jax.random.key(0, impl="philox2x32"))]:
    print(f"{name:28s} {str(kk.shape):>10s} {str(jax.random.key_data(kk).shape):>15s}")

assert jax.random.split(k, (2, 3)).shape == (2, 3)
assert jax.random.key_data(jax.random.split(k, (2, 3))).shape == (2, 3, 2)
assert jax.random.key_data(jax.random.split(jax.random.key(0, impl="rbg"), (2, 3))).shape == (2, 3, 4)
assert jax.random.key_data(jax.random.key(0, impl="philox2x32")).shape == (1,)

ks = jax.random.split(k, 4)
ok_ops = {
    "ks[1:3]（切片）":            lambda: ks[1:3].shape,
    "ks.reshape(2, 2)":           lambda: ks.reshape(2, 2).shape,
    "ks.reshape(2, 2).T":         lambda: ks.reshape(2, 2).T.shape,
    "jnp.stack([k, k])":          lambda: jnp.stack([k, k]).shape,
    "jnp.concatenate([ks, ks])":  lambda: jnp.concatenate([ks, ks]).shape,
    "jnp.where(mask, ks, ks[::-1])": lambda: jnp.where(jnp.array([True, False, True, False]), ks, ks[::-1]).shape,
    "jnp.broadcast_to(k, (3,))":  lambda: jnp.broadcast_to(k, (3,)).shape,
    "ks == ks（比较）":           lambda: (ks == ks).tolist(),
}
for name, f in ok_ops.items():
    print(f"✅ {name:32s} → {f()}")
assert jnp.stack([k, k]).dtype == k.dtype and bool(jnp.all(ks == ks))

# ❌ 错误写法：对类型化 key 做算术 / 规约 / 类型转换（想“手工派生”新 key 或把 key 当数据保存）
bad_ops = {
    "k + 1":                     lambda: k + 1,
    "k * 2":                     lambda: k * 2,
    "k ^ k":                     lambda: k ^ k,
    "jnp.sum(ks)":               lambda: jnp.sum(ks),
    "jnp.unique(ks)":            lambda: jnp.unique(ks),
    "jnp.asarray(k, jnp.uint32)": lambda: jnp.asarray(k, jnp.uint32),
    "np.asarray(k)":             lambda: np.asarray(k),
}
bad_errors = {}
for name, f in bad_ops.items():
    try:
        out = f()
        print(f"（没有报错）{name} → {out}")
    except Exception as e:
        bad_errors[name] = type(e).__name__
        print(f"❌ {name:28s} → {err_line(e)}")

assert set(bad_errors) == set(bad_ops)          # 全部报错
assert bad_errors["k + 1"] == "TypeError" and bad_errors["np.asarray(k)"] == "TypeError"
# 对照：原始 uint32 key 做同样的事会“静默成功”，这正是旧写法的隐患
print("\n原始 key：PRNGKey(0) + 1 =", raw + 1, "（静默成功）")
assert jnp.array_equal(raw + 1, jnp.array([1, 1], jnp.uint32))

# ✅ 正确写法：派生新 key 用 fold_in / split；取底层比特用 key_data
k_step1 = jax.random.fold_in(k, 1)                  # 代替 “k + 1”
raw_bits = jax.random.key_data(k)                   # 代替 “np.asarray(k)”
print("fold_in(k, 1) 的底层数据:", jax.random.key_data(k_step1))
print("key_data(k) =", raw_bits, raw_bits.dtype)
assert k_step1.dtype == k.dtype and not bool(k_step1 == k)
assert raw_bits.dtype == jnp.uint32 and raw_bits.shape == (2,)

# ❌ 错误写法：只保存了 rbg key 的原始数据，恢复时没有指定 impl（按默认的 threefry2x32 解释）
rbg_key = jax.random.key(7, impl="rbg")
rbg_raw = jax.random.key_data(rbg_key)
try:
    restored = jax.random.wrap_key_data(rbg_raw)
    print("（没有报错）恢复为", restored.dtype)
    wrap_err = None
except Exception as e:
    wrap_err = e
    print("❌", err_line(e))
assert isinstance(wrap_err, TypeError)

# ✅ 正确写法：保存 (impl 名称, key_data)，恢复时 wrap_key_data(..., impl=impl)
def save_keys(keys):
    """把一个类型化 key 数组序列化成 bytes：impl 名称 + 原始 uint32 数据。"""
    buf = io.BytesIO()
    np.savez(buf, impl=str(jax.random.key_impl(keys)), data=np.asarray(jax.random.key_data(keys)))
    return buf.getvalue()

def load_keys(blob):
    z = np.load(io.BytesIO(blob))
    return jax.random.wrap_key_data(jnp.asarray(z["data"]), impl=str(z["impl"]))

for impl in ["threefry2x32", "rbg", "philox2x32"]:
    keys = jax.random.split(jax.random.key(7, impl=impl), (2, 3))
    back = load_keys(save_keys(keys))
    same = bool(jnp.all(back == keys)) and back.dtype == keys.dtype
    same_samples = jnp.array_equal(jax.random.normal(back[1, 2], (4,)), jax.random.normal(keys[1, 2], (4,)))
    print(f"{impl:13s} 往返后 dtype={back.dtype}, shape={back.shape}, key 相同={same}, 采样相同={bool(same_samples)}")
    assert same and same_samples

# (1) 种子截断
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    same_key = bool(jax.random.key(2**32 + 5) == jax.random.key(5))
print("32 位模式：key(2**32 + 5) == key(5)?", same_key, "| 警告数:", len(caught))
assert same_key and len(caught) == 0

x64_res = run_py("""
import json, jax
jax.config.update("jax_enable_x64", True)
a = jax.random.key(2**32 + 5)
print("JSON:" + json.dumps({"data": jax.random.key_data(a).tolist(), "same": bool(a == jax.random.key(5))}))
""")
print("x64 模式（子进程）：key_data(key(2**32 + 5)) =", x64_res["data"], "| 与 key(5) 相同?", x64_res["same"])
assert x64_res["data"] == [1, 5] and not x64_res["same"]

# 大整数种子的稳妥写法：拆成高低两段依次 fold_in，结果与 x64 设置无关
def key_from_big_seed(seed: int):
    hi, lo = (seed >> 32) & 0xFFFFFFFF, seed & 0xFFFFFFFF
    return jax.random.fold_in(jax.random.fold_in(jax.random.key(0), hi), lo)
assert not bool(key_from_big_seed(2**32 + 5) == key_from_big_seed(5))

# (2) jax_legacy_prng_key
print("\njax_legacy_prng_key 默认值:", jax.config.jax_legacy_prng_key)
with jax.legacy_prng_key("error"):
    try:
        jax.random.normal(jax.random.PRNGKey(0), (2,))
        legacy_err = None
    except Exception as e:
        legacy_err = e
        print("'error' 模式下传入原始 key →", err_line(e))
    ok_typed = jax.random.normal(jax.random.key(0), (2,))      # 类型化 key 不受影响
assert isinstance(legacy_err, ValueError) and ok_typed.shape == (2,)

IMPLS = ["threefry2x32", "threefry4x32", "philox2x32", "philox4x32", "rbg", "unsafe_rbg"]
samples = {}
print(f"{'impl':13s} {'dtype':10s} {'key_data':>9s}  normal(key(0), (3,))")
for impl in IMPLS:
    kk = jax.random.key(0, impl=impl)
    samples[impl] = np.asarray(jax.random.normal(kk, (3,)))
    print(f"{impl:13s} {str(kk.dtype):10s} {str(jax.random.key_data(kk).shape):>9s}  {samples[impl]}")

expected_dtype = {"threefry2x32": "key<fry>", "threefry4x32": "key<fry4>", "philox2x32": "key<phx2>",
                  "philox4x32": "key<phx4>", "rbg": "key<rbg>", "unsafe_rbg": "key<urbg>"}
assert all(str(jax.random.key(0, impl=i).dtype) == d for i, d in expected_dtype.items())
# 同一个种子，不同算法给出不同的随机数（rbg 与 unsafe_rbg 的“生成”部分相同，只有派生 key 不同）
distinct = {tuple(v.tolist()) for v in samples.values()}
assert len(distinct) == 5 and np.array_equal(samples["rbg"], samples["unsafe_rbg"])

# 全局默认算法：环境变量（子进程，从启动时生效）与上下文管理器（本进程，临时）
env_res = run_py("""
import json, jax
print("JSON:" + json.dumps({
    "config": jax.config.jax_default_prng_impl,
    "key_dtype": str(jax.random.key(0).dtype),
    "raw_shape": list(jax.random.PRNGKey(0).shape),                      # 原始 key 按全局配置解释
    "explicit": str(jax.random.key(0, impl="threefry2x32").dtype)}))     # 显式 impl 不受影响
""", env={"JAX_DEFAULT_PRNG_IMPL": "rbg"})
print("JAX_DEFAULT_PRNG_IMPL=rbg 的子进程:", env_res)
assert env_res == {"config": "rbg", "key_dtype": "key<rbg>", "raw_shape": [4], "explicit": "key<fry>"}

with jax.default_prng_impl("philox4x32"):
    inside = str(jax.random.key(0).dtype)
outside = str(jax.random.key(0).dtype)
print("with jax.default_prng_impl('philox4x32') 内:", inside, "| 退出后:", outside)
assert inside == "key<phx4>" and outside == "key<fry>"

print(f"{'impl':13s} {'生成 4M normal':>14s} {'split 10万':>11s} {'首次编译':>9s}")
for impl in IMPLS:
    kk = jax.random.key(0, impl=impl)
    gen = jax.jit(lambda key: jax.random.normal(key, (4_000_000,)))
    spl = jax.jit(lambda key: jax.random.split(key, 100_000))
    t0 = time.perf_counter(); gen.lower(kk).compile(); t_compile = (time.perf_counter() - t0) * 1e3
    print(f"{impl:13s} {bench(gen, kk):11.2f} ms {bench(spl, kk):8.2f} ms {t_compile:6.0f} ms")

kr = jax.random.key(0, impl="rbg")
ku = jax.random.key(0, impl="unsafe_rbg")
halves = jax.random.key_data(kr).reshape(2, 2)          # rbg key = 两个 threefry key 拼起来
tf_split = jnp.concatenate([jax.random.key_data(jax.random.split(jax.random.wrap_key_data(h), 3)) for h in halves], -1)

rbg_like_tf = bool(jnp.array_equal(jax.random.key_data(jax.random.split(kr, 3)), tf_split))
urbg_like_tf = bool(jnp.array_equal(jax.random.key_data(jax.random.split(ku, 3)), tf_split))
print("rbg 的 split == 两半各自做 threefry split:", rbg_like_tf)
print("unsafe_rbg 的 split == 同上:", urbg_like_tf, "（它用 RBG 当哈希，质量未经检验）")
assert rbg_like_tf and not urbg_like_tf

# ❌ 错误写法：用 rbg 时默认“vmap 等于逐个调用”（例如用 vmap 版和循环版对拍）
keys_rbg = jax.random.split(jax.random.key(0, impl="rbg"), 8)
try:
    v = jax.vmap(lambda kk: jax.random.normal(kk))(keys_rbg)
    loop = jnp.stack([jax.random.normal(kk) for kk in keys_rbg])
    first_only = jax.random.normal(keys_rbg[0], (8,))
    print("没有报错，但 vmap == 逐个调用?", bool(jnp.array_equal(v, loop)),
          "| vmap == normal(keys[0], (8,))?", bool(jnp.array_equal(v, first_only)))
except Exception as e:
    print("❌", err_line(e))
assert not jnp.array_equal(v, loop) and jnp.array_equal(v, first_only)

# ✅ 正确写法：需要“vmap 与逐个调用一致”时用 threefry / philox（默认的 threefry2x32 即可）
for impl in ["threefry2x32", "philox4x32"]:
    keys_ok = jax.random.split(jax.random.key(0, impl=impl), 8)
    v_ok = jax.vmap(lambda kk: jax.random.normal(kk))(keys_ok)
    loop_ok = jnp.stack([jax.random.normal(kk) for kk in keys_ok])
    print(f"{impl:13s} vmap == 逐个调用: {bool(jnp.array_equal(v_ok, loop_ok))}")
    assert jnp.array_equal(v_ok, loop_ok)

print("jax_threefry_partitionable =", jax.config.jax_threefry_partitionable)
assert jax.config.jax_threefry_partitionable is True
# jax.random 的模块文档写明了从哪个版本起成为默认
assert "As of JAX v.0.5.0, this is the" in jax.random.__doc__
print("jax.random 文档：'As of JAX v.0.5.0, this is the default.' ✓")

# 三种修改方式：jax.config.update / 环境变量 JAX_THREEFRY_PARTITIONABLE / 上下文管理器
env_flag = run_py('import json, jax; print("JSON:" + json.dumps(jax.config.jax_threefry_partitionable))',
                  env={"JAX_THREEFRY_PARTITIONABLE": "0"})
with jax.threefry_partitionable(False):
    ctx_flag = jax.config.jax_threefry_partitionable
print("环境变量 JAX_THREEFRY_PARTITIONABLE=0 →", env_flag, "| 上下文管理器内 →", ctx_flag,
      "| 退出后 →", jax.config.jax_threefry_partitionable)
assert env_flag is False and ctx_flag is False and jax.config.jax_threefry_partitionable is True

FLAG_SCRIPT = """
import sys, json, jax, jax.numpy as jnp
jax.config.update("jax_threefry_partitionable", sys.argv[1] == "1")
k = jax.random.key(0)
x10 = jax.random.normal(k, (10,))
print("JSON:" + json.dumps({
    "flag": jax.config.jax_threefry_partitionable,
    "normal_8": jax.random.normal(k, (8,)).tolist(),
    "uniform_3": jax.random.uniform(k, (3,)).tolist(),
    "split": jax.random.key_data(jax.random.split(k)).tolist(),
    "fold_in": jax.random.key_data(jax.random.fold_in(k, 1)).tolist(),
    "dot10": float(jnp.dot(x10, x10)),      # 第 1 章 1.4 节安装验证脚本的那个数
}))
"""
flag_runs = {flag: [run_py(FLAG_SCRIPT, int(flag)) for _ in range(2)] for flag in (True, False)}
on, off = flag_runs[True][0], flag_runs[False][0]
for name in ["normal_8", "uniform_3", "split", "fold_in", "dot10"]:
    print(f"{name:10s} 开: {on[name]}\n{'':10s} 关: {off[name]}\n")

assert on["flag"] is True and off["flag"] is False
assert flag_runs[True][0] == flag_runs[True][1] and flag_runs[False][0] == flag_runs[False][1]   # 同一设置下可复现
assert on["normal_8"] != off["normal_8"] and on["uniform_3"] != off["uniform_3"]               # 随机比特不同
assert on["split"] != off["split"]                                                              # split 也不同
assert on["fold_in"] == off["fold_in"]                                                          # fold_in 不受影响
# 与第 1 章 1.4 节的实测值一致：开 9.7387…，关 2.6585…
assert np.isclose(on["dot10"], 9.738718032836914) and np.isclose(off["dot10"], 2.6585559844970703)
# 本进程的默认设置与“开”的子进程逐位相同
assert jax.random.normal(jax.random.key(0), (8,)).tolist() == on["normal_8"]
print("✓ 开 / 关 两种设置下：随机比特和 split 都不同，fold_in 相同；各自可复现")

from jax.extend.random import threefry_2x32, threefry2x32_p

kd = jax.random.key_data(k)
n = 8
with jax.threefry_partitionable(False):
    bits_off = jax.random.bits(k, (n,))
    manual_off = threefry_2x32(kd, jnp.arange(n, dtype=jnp.uint32))   # 内部把计数器劈成两半
with jax.threefry_partitionable(True):
    bits_on = jax.random.bits(k, (n,))
    w0, w1 = threefry2x32_p.bind(kd[0], kd[1], jnp.zeros(n, jnp.uint32), jnp.arange(n, dtype=jnp.uint32))
    manual_on = w0 ^ w1                                                  # 计数器 = 全局下标 i，两字异或
print("关：bits      =", bits_off)
print("关：手工复现  =", manual_off)
print("开：bits      =", bits_on)
print("开：手工复现  =", manual_on)
assert jnp.array_equal(bits_off, manual_off) and jnp.array_equal(bits_on, manual_on)

for flag in (True, False):
    with jax.threefry_partitionable(flag):
        prefix_1d = bool(jnp.array_equal(jax.random.normal(k, (10,))[:5], jax.random.normal(k, (5,))))
        prefix_2d = bool(jnp.array_equal(jax.random.normal(k, (4, 5))[:2], jax.random.normal(k, (2, 5))))
        split_prefix = bool(jnp.all(jax.random.split(k, 5)[:3] == jax.random.split(k, 3)))
        split_is_fold = bool(jnp.all(jax.random.split(k, 4) == jnp.stack([jax.random.fold_in(k, i) for i in range(4)])))
    print(f"开关={flag!s:5}: 前缀性质 1D={prefix_1d}, 2D 行前缀={prefix_2d}, "
          f"split(k,5)[:3]==split(k,3): {split_prefix}, split(k,n)[i]==fold_in(k,i): {split_is_fold}")
    assert prefix_1d == prefix_2d == split_prefix == split_is_fold == flag

mesh = jax.make_mesh((8,), ("x",))
N = 1024
k_sh = jax.random.key(0)

def threefry_values_per_device(hlo_text, flag):
    """从 HLO 中 threefry 轮函数的 xor 形状，估算每个设备生成的随机值个数。"""
    shapes = re.findall(r"u32\[([0-9,]+)\][^\n]*? xor\(", hlo_text)
    elems = max(int(np.prod([int(d) for d in s.split(",")])) for s in shapes)
    return elems * (1 if flag else 2)

rows = []
for flag in (True, False):
    with jax.threefry_partitionable(flag):
        ref = jax.jit(lambda key: jax.random.normal(key, (8, N)))(k_sh)          # 不分片的参考结果
        for spec in (P("x"), P(None, "x")):
            f = jax.jit(lambda key: jax.random.normal(key, (8, N)), out_shardings=NamedSharding(mesh, spec))
            out = f(k_sh)
            txt = f.lower(k_sh).compile().as_text()
            rows.append(dict(flag=flag, spec=str(spec), equal=bool(jnp.array_equal(out, ref)),
                             comm=collectives(txt), per_dev=threefry_values_per_device(txt, flag),
                             shard_shape=out.addressable_shards[0].data.shape))

print(f"{'开关':5s} {'分片':14s} {'与不分片相同':>10s} {'每设备分片':>10s} {'每设备生成':>9s}  集合通信")
for r in rows:
    comm = {kk: v for kk, v in r["comm"].items() if v} or "无"
    print(f"{r['flag']!s:5s} {r['spec']:14s} {r['equal']!s:>10s} {str(r['shard_shape']):>12s} {r['per_dev']:>9d}  {comm}")

R = {(r["flag"], r["spec"]): r for r in rows}
assert all(r["equal"] for r in rows)                                     # 4 种情况结果都与不分片一致
# 开启：两种切法都没有任何集合通信，每个设备只生成自己那 8*N/8 个值
assert all(sum(R[(True, s)]["comm"].values()) == 0 and R[(True, s)]["per_dev"] == N for s in ("P('x',)", "P(None, 'x')"))
# 关闭 + 按行切：需要集合通信（all-gather / collective-permute）把成对的元素换到正确的设备
assert sum(R[(False, "P('x',)")]["comm"].values()) > 0
# 关闭 + 按列切：没有通信，但每个设备都把整个 8*N 数组算了一遍（8 倍冗余计算）
assert sum(R[(False, "P(None, 'x')")]["comm"].values()) == 0 and R[(False, "P(None, 'x')")]["per_dev"] == 8 * N
print("\n✓ 开启：本地生成、零通信；关闭：要么通信，要么每个设备冗余地算整个数组")

SEARCH_SCRIPT = """
import sys, json, jax, jax.numpy as jnp, numpy as np
from jax import lax
jax.config.update("jax_threefry_partitionable", sys.argv[1] == "1")
jax.config.update("jax_enable_x64", sys.argv[2] == "1")
TARGET = 10.811179161071777
out = {}
for impl in ["threefry2x32", "threefry4x32", "philox2x32", "philox4x32", "rbg", "unsafe_rbg"]:
    def one(seed):
        x = jax.random.normal(jax.random.key(seed, impl=impl), (10,))
        return jnp.dot(x, x)
    vals = np.asarray(jax.jit(lambda s: lax.map(one, s))(jnp.arange(1000)), np.float64)
    d = np.abs(vals - TARGET)
    out[impl] = {"seed0": float(vals[0]), "min_dist": float(d.min()), "argmin": int(d.argmin())}
print("JSON:" + json.dumps(out))
"""
search = {(f, x): run_py(SEARCH_SCRIPT, f, x) for f in (1, 0) for x in (0, 1)}
best = min(((v["min_dist"], cfg, impl, v["argmin"]) for cfg, res in search.items() for impl, v in res.items()))
print(f"{'partitionable, x64':20s} {'threefry2x32 种子 0 的 dot(x,x)':>32s}")
for (f, x), res in search.items():
    print(f"{f!s:>10s}, {x!s:<8s} {res['threefry2x32']['seed0']:>32.6f}")
print(f"\n全部 {len(search) * 6 * 1000} 个组合中离 10.811179 最近的: 距离 {best[0]:.2e}"
      f"（partitionable={best[1][0]}, x64={best[1][1]}, impl={best[2]}, seed={best[3]}）")

assert np.isclose(search[(1, 0)]["threefry2x32"]["seed0"], 9.7387, atol=1e-3)     # 0.12 默认
assert np.isclose(search[(0, 0)]["threefry2x32"]["seed0"], 2.6586, atol=1e-3)     # 0.5.0 之前的默认
assert best[0] > 1e-4          # 没有任何组合能在浮点误差内复现 10.811179（float32 误差约 1e-6）

# ❌ 错误写法：split 后没有更新 key，下一次又从同一个旧 key 拆分 → 两次拿到同一个子 key
key = jax.random.key(0)
try:
    _, k_init = jax.random.split(key)          # 新 key 被丢进了 _
    w_init = jax.random.normal(k_init, (4,))
    _, k_noise = jax.random.split(key)         # 又从旧 key 拆分
    noise = jax.random.normal(k_noise, (4,))
    print("没有报错，但“初始化”和“噪声”完全相同:", bool(jnp.array_equal(w_init, noise)))
except Exception as e:
    print("❌", err_line(e))
assert jnp.array_equal(w_init, noise)

# ✅ 正确写法：每次 split 都把新 key 赋回给 key，旧 key 不再使用
key = jax.random.key(0)
key, k_init = jax.random.split(key)
w_init = jax.random.normal(k_init, (4,))
key, k_noise = jax.random.split(key)
noise = jax.random.normal(k_noise, (4,))
print("初始化与噪声不同:", not bool(jnp.array_equal(w_init, noise)))
assert not jnp.array_equal(w_init, noise) and not bool(k_init == k_noise)

# ❌ 错误写法：用 key(seed + i) 给第 i 个 worker 派生 key → 不同种子的实验之间共享随机数
def worker_keys_bad(seed, n):
    return [jax.random.key(seed + i) for i in range(n)]
try:
    run_seed0 = worker_keys_bad(0, 4)
    run_seed1 = worker_keys_bad(1, 4)
    collide = bool(run_seed0[1] == run_seed1[0])
    print("没有报错，但 种子0/worker1 与 种子1/worker0 的 key 相同:", collide)
except Exception as e:
    print("❌", err_line(e))
assert collide
assert jnp.array_equal(jax.random.normal(run_seed0[1], (3,)), jax.random.normal(run_seed1[0], (3,)))

# ✅ 正确写法：从实验的主 key 一次 split 出 n 个
def worker_keys(seed, n):
    return jax.random.split(jax.random.key(seed), n)

a, b = worker_keys(0, 4), worker_keys(1, 4)
all_keys = [tuple(r) for r in np.asarray(jax.random.key_data(jnp.concatenate([a, b]))).tolist()]
print("两次实验共 8 个 key，互不相同:", len(set(all_keys)) == 8)
assert len(set(all_keys)) == 8

# 一次 split(key, n) vs Python 循环里 n 次 split（结果都可用，但循环要 n 次派发）
def loop_split(key, n):
    out = []
    for _ in range(n):
        key, sub = jax.random.split(key)
        out.append(sub)
    return out
t_loop = bench(lambda kk: loop_split(kk, 200), jax.random.key(0), runs=5)      # bench 先预热，排除首次编译
t_once = bench(lambda kk: jax.random.split(kk, 200), jax.random.key(0))
print(f"200 个 key：Python 循环 split {t_loop:.2f} ms，一次 split(key, 200) {t_once:.3f} ms（预热后中位数，只打印）")

# ❌ 错误写法：训练循环每一步都用同一个 key 生成 dropout 掩码 → 每步掩码都一样
base = jax.random.key(42)
x_act = jnp.ones((4, 6))

@jax.jit
def dropout_mask(key, x):
    return jax.random.bernoulli(key, 0.5, x.shape)

try:
    masks_bad = [dropout_mask(base, x_act) for step in range(5)]
    print("没有报错，但 5 步的掩码全部相同:", all(bool(jnp.array_equal(m, masks_bad[0])) for m in masks_bad))
except Exception as e:
    print("❌", err_line(e))
assert all(jnp.array_equal(m, masks_bad[0]) for m in masks_bad)

# ✅ 正确写法：在 jit 内用 fold_in(base, step) 派生每一步的 key，step 作为动态参数传入
trace_count = 0
@jax.jit
def dropout_mask_at(base, step, x):
    global trace_count
    trace_count += 1
    return jax.random.bernoulli(jax.random.fold_in(base, step), 0.5, x.shape)

masks = [dropout_mask_at(base, jnp.uint32(s), x_act) for s in range(5)]
distinct_masks = len({np.asarray(m).tobytes() for m in masks})
print("5 步的掩码互不相同:", distinct_masks == 5, "| 追踪次数:", trace_count)
assert distinct_masks == 5 and trace_count == 1
# 随机访问：直接重算第 3 步，不用重放前面的步
assert jnp.array_equal(dropout_mask_at(base, jnp.uint32(3), x_act), masks[3])

# 按数据下标派生：同一个样本无论被分进哪个 batch，增强噪声都相同
k_aug = jax.random.key(7)
def augment(idx, x):                          # x 是单个样本，idx 是它在数据集中的下标
    return x + 0.1 * jax.random.normal(jax.random.fold_in(k_aug, idx), x.shape)
data = jnp.arange(10.0).reshape(10, 1) * jnp.ones((10, 3))
batch_a = jnp.array([5, 1, 7]); batch_b = jnp.array([2, 5, 9, 0])
aug_a = jax.vmap(augment)(batch_a, data[batch_a])
aug_b = jax.vmap(augment)(batch_b, data[batch_b])
print("样本 5 在两个不同 batch 中的增强结果相同:", bool(jnp.array_equal(aug_a[0], aug_b[1])))
assert jnp.array_equal(aug_a[0], aug_b[1])

try:
    jax.random.fold_in(base, -1)
except OverflowError as e:
    print("fold_in(base, -1) →", err_line(e))

# ❌ 错误写法：key 放进了 carry，但循环体没有更新它 → 每一步的随机数都一样
def body_bad(key, _):
    return key, jax.random.normal(key)             # 原样把旧 key 传给下一步

try:
    _, ys_bad = lax.scan(body_bad, jax.random.key(0), None, length=6)
    print("没有报错，但 6 步的随机数:", ys_bad)
except Exception as e:
    print("❌", err_line(e))
assert bool(jnp.all(ys_bad == ys_bad[0]))

# ✅ 正确写法 A：每一步 split，把新 key 作为 carry 交给下一步
def body(key, _):
    key, sub = jax.random.split(key)
    return key, jax.random.normal(sub)
_, ys_a = jax.jit(lambda kk: lax.scan(body, kk, None, length=6))(jax.random.key(0))

# 与等价的 Python 循环逐位相同
kk, ys_loop = jax.random.key(0), []
for _ in range(6):
    kk, sub = jax.random.split(kk)
    ys_loop.append(jax.random.normal(sub))
assert jnp.array_equal(ys_a, jnp.stack(ys_loop))

# ✅ 正确写法 B：事先 split 出 T 个 key，作为 xs 逐步喂入（carry 里不放 key）
_, ys_b = lax.scan(lambda c, kt: (c, jax.random.normal(kt)), None, jax.random.split(jax.random.key(0), 6))
print("写法 A:", ys_a)
print("写法 B:", ys_b)
assert len(set(np.asarray(ys_a).tolist())) == 6 and len(set(np.asarray(ys_b).tolist())) == 6

# ❌ 错误写法：vmap 时 key 不映射（in_axes=None）→ 每个样本拿到同一组“随机数”
def noisy_sample(key, x):
    return x + jax.random.normal(key, x.shape)

xs_batch = jnp.zeros((4, 3))
try:
    out_bad = jax.vmap(noisy_sample, in_axes=(None, 0))(jax.random.key(0), xs_batch)
    print("没有报错，但 4 个样本的噪声完全相同:", bool(jnp.all(out_bad == out_bad[0])))
except Exception as e:
    print("❌", err_line(e))
assert bool(jnp.all(out_bad == out_bad[0]))

# ✅ 正确写法：split 出 B 个 key，与样本一起沿第 0 轴映射
keys_b = jax.random.split(jax.random.key(0), xs_batch.shape[0])
out_ok = jax.vmap(noisy_sample)(keys_b, xs_batch)
loop_ok = jnp.stack([noisy_sample(kb, xb) for kb, xb in zip(keys_b, xs_batch)])
print("vmap 结果与逐个调用相同:", bool(jnp.array_equal(out_ok, loop_ok)))
assert jnp.array_equal(out_ok, loop_ok)
assert len({np.asarray(r).tobytes() for r in out_ok}) == 4

# 两层 vmap：split(key, (B, T)) 直接得到二维 key 数组
keys_bt = jax.random.split(jax.random.key(1), (4, 5))
out_bt = jax.vmap(jax.vmap(lambda kk: jax.random.normal(kk)))(keys_bt)
assert out_bt.shape == (4, 5) and out_bt[2, 3] == jax.random.normal(keys_bt[2, 3])

# ❌ 错误写法：直接用进程编号当种子（这里用参数模拟 process_index）
def process_key_bad(seed, process_index):
    return jax.random.key(process_index)          # 种子被忽略了！
try:
    collide = bool(process_key_bad(0, 1) == process_key_bad(123, 1))
    print("没有报错，但 种子0 和 种子123 两次实验里 进程1 的 key 相同:", collide)
except Exception as e:
    print("❌", err_line(e))
assert collide

# ✅ 正确写法：实验种子决定主 key，再 fold_in 进程编号
def process_key(seed, process_index):
    return jax.random.fold_in(jax.random.key(seed), process_index)
pkeys = [process_key(s, p) for s in (0, 123) for p in range(4)]
print("2 次实验 × 4 个进程 = 8 个 key，互不相同:",
      len({tuple(np.asarray(jax.random.key_data(x)).tolist()) for x in pkeys}) == 8)
assert len({tuple(np.asarray(jax.random.key_data(x)).tolist()) for x in pkeys}) == 8
assert bool(process_key(0, jax.process_index()) == jax.random.fold_in(jax.random.key(0), 0))   # 单进程时 index=0

MP_WORKER = """
import os, sys, json
pid, nproc, port = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3]
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"   # 每个进程 4 个设备
import jax, jax.numpy as jnp, numpy as np
from jax.sharding import NamedSharding, PartitionSpec as P
jax.distributed.initialize(coordinator_address=f"localhost:{port}", num_processes=nproc, process_id=pid)
mesh = jax.make_mesh((jax.device_count(),), ("data",))
gen = jax.jit(lambda k: jax.random.normal(k, (8, 4)), out_shardings=NamedSharding(mesh, P("data")))
base = jax.random.key(42)
def local_rows(arr):      # 本进程能访问的分片：{起始行号: 数据}
    return {int(s.index[0].start): np.asarray(s.data).tolist() for s in arr.addressable_shards}
out = {"pid": jax.process_index(), "nproc": jax.process_count(), "ndev": jax.device_count(),
       "shared": local_rows(gen(base)),                                          # ✅ 同一个 key
       "per_process": local_rows(gen(jax.random.fold_in(base, jax.process_index()))),   # ❌ 各进程不同的 key
       "local_perm": np.asarray(jax.random.permutation(jax.random.fold_in(base, jax.process_index()), 6)).tolist()}
print("JSON:" + json.dumps(out))
"""

def run_two_processes(code, nproc=2, timeout=300):
    s = socket.socket(); s.bind(("localhost", 0)); port = s.getsockname()[1]; s.close()
    procs = [subprocess.Popen([sys.executable, "-c", code, str(i), str(nproc), str(port)],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) for i in range(nproc)]
    results = {}
    for p in procs:
        o, e = p.communicate(timeout=timeout)
        assert p.returncode == 0, e[-3000:]
        r = json.loads([l for l in o.splitlines() if l.startswith("JSON:")][-1][5:])
        results[r["pid"]] = r
    return results

mp = run_two_processes(MP_WORKER)

def assemble(field):
    rows = {}
    for r in mp.values():
        rows.update({int(i): v for i, v in r[field].items()})
    return np.concatenate([np.array(rows[i]) for i in sorted(rows)])

ref42 = np.asarray(jax.random.normal(jax.random.key(42), (8, 4)))
ref_p0 = np.asarray(jax.random.normal(jax.random.fold_in(jax.random.key(42), 0), (8, 4)))
ref_p1 = np.asarray(jax.random.normal(jax.random.fold_in(jax.random.key(42), 1), (8, 4)))
g_shared, g_bad = assemble("shared"), assemble("per_process")

for pid, r in sorted(mp.items()):
    print(f"进程 {pid}/{r['nproc']}：全局设备 {r['ndev']} 个，本进程持有的行 {sorted(int(i) for i in r['shared'])}，本地打乱 {r['local_perm']}")
print("✅ 同一个 key：拼出的全局数组 == 单进程 normal(key(42), (8,4))?", np.array_equal(g_shared, ref42))
print("❌ 各进程不同 key：前 4 行 == key A 的前 4 行?", np.array_equal(g_bad[:4], ref_p0[:4]),
      "| 后 4 行 == key B 的后 4 行?", np.array_equal(g_bad[4:], ref_p1[4:]),
      "| 等于任何一个完整数组?", any(np.array_equal(g_bad, r_) for r_ in (ref42, ref_p0, ref_p1)))

assert all(r["nproc"] == 2 and r["ndev"] == 8 for r in mp.values())
assert np.array_equal(g_shared, ref42)
assert np.array_equal(g_bad[:4], ref_p0[:4]) and np.array_equal(g_bad[4:], ref_p1[4:])
assert not any(np.array_equal(g_bad, r_) for r_ in (ref42, ref_p0, ref_p1))
assert mp[0]["local_perm"] != mp[1]["local_perm"]