第 3 章　JAX 分布式计算

# 必须在 import jax 之前设置：用 1 个 CPU 模拟 8 个设备（要求重启内核后执行）
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import re, sys, json, subprocess, tempfile, socket, warnings
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P, AxisType
import optax

assert jax.device_count() == 8, "设备数不是 8：请重启内核后从第一个单元开始执行"
N = jax.device_count()
print("JAX", jax.__version__, "| 设备数:", N, "| 平台:", jax.devices()[0].platform)


def f64(a):
    """转成 NumPy float64，用来和分布式结果对拍"""
    return np.asarray(a, dtype=np.float64)


def collectives(compiled):
    """从编译后的 HLO 文本中找出用到的集合通信算子"""
    return sorted(set(re.findall(r"\b(all-reduce|reduce-scatter|all-gather|all-to-all|collective-permute)\(",
                                 compiled.as_text())))

# 验证上面的公式：N 份等大分片上「本地平均损失的梯度」取平均 == 全局批次的梯度（NumPy float64 对拍）
rng = np.random.default_rng(0)
Xg = rng.normal(size=(64, 5)); yg = rng.normal(size=64); w0 = rng.normal(size=5)


def grad_mse(w, X, y):                      # 线性回归 MSE 的解析梯度：2/|B| * X^T (Xw - y)
    return 2.0 / len(X) * X.T @ (X @ w - y)


g_global = grad_mse(w0, Xg, yg)
g_local = [grad_mse(w0, Xs, ys) for Xs, ys in zip(np.split(Xg, 8), np.split(yg, 8))]
assert np.allclose(np.mean(g_local, axis=0), g_global)          # pmean 本地梯度 == 全局梯度
assert np.allclose(np.sum(g_local, axis=0), 8 * g_global)       # psum 本地梯度 == N 倍全局梯度
print("pmean(本地梯度) == 全局梯度 ✓；psum(本地梯度) == 8 × 全局梯度 ✓")

# 流水线气泡比例 (p-1)/(m+p-1)：p=4 个阶段时，微批次越多气泡越小
def bubble(p, m):
    return (p - 1) / (m + p - 1)


for m in [1, 4, 16, 64]:
    print(f"p=4, m={m:3d}: 气泡比例 = {bubble(4, m):.3f}")
assert bubble(4, 1) == 0.75 and bubble(4, 64) < 0.05   # m=1 即朴素模型并行，3/4 的时间在等待

# 课程示例：查看设备
print("jax.devices():          ", jax.devices())
print("jax.device_count():     ", jax.device_count())
print("jax.local_devices():    ", jax.local_devices()[:3], "...")
print("jax.local_device_count():", jax.local_device_count())

d0 = jax.local_devices()[0]
print(f"第一个本地设备: ID={d0.id}, Platform={d0.platform}, process_index={d0.process_index}, kind={d0.device_kind}")

assert jax.local_devices() == jax.devices()          # 单进程：本地 == 全局
assert all(d.process_index == 0 for d in jax.devices())
assert [d.id for d in jax.devices()] == list(range(8))

# ❌ 旧写法（过时）：x.device()。在 0.12 中 device 是属性，不是方法
x = jnp.ones((3, 3))
try:
    x.device()
    raised = False
except TypeError as e:
    raised = True
    print("x.device() 报错:", type(e).__name__, "-", e)
assert raised

# ✅ 推荐写法：x.device 属性 / x.devices() / x.sharding
x = jnp.ones((3, 3))
print("x.device   :", x.device)
print("x.devices():", x.devices())
print("x.sharding :", x.sharding)
assert x.device == jax.devices()[0]
assert x.devices() == {jax.devices()[0]}

# 课程示例：本地设备多于 1 个时才演示显式放置
num_local_devices = jax.local_device_count()
if num_local_devices > 1:
    host_array = np.random.default_rng(0).random((2, 2))
    print("放置前类型:", type(host_array))
    target = jax.local_devices()[1]
    dev_array = jax.device_put(host_array, target)
    print("放置后类型:", type(dev_array).__name__, "| 所在设备:", dev_array.device)
    assert dev_array.device == target
elif num_local_devices == 1:
    print("只有一个本地设备可用，跳过明确放置示例。")
else:
    print("未找到 JAX 设备。")

u = jnp.ones(3)                                        # uncommitted（默认在 cpu:0）
c1 = jax.device_put(jnp.ones(3), jax.devices()[1])     # committed 到 cpu:1
c2 = jax.device_put(jnp.ones(3), jax.devices()[2])     # committed 到 cpu:2

r = u + c1
print("uncommitted + committed(cpu:1) → 结果在", r.device)
assert r.device == jax.devices()[1]

try:
    c1 + c2
    raised = False
except ValueError as e:
    raised = True
    print("committed(cpu:1) + committed(cpu:2) 报错:", str(e)[:110], "...")
assert raised

# 需要跨设备运算时，要显式搬运
r2 = c1 + jax.device_put(c2, jax.devices()[1])
assert r2.device == jax.devices()[1]

# 在干净的子进程中演示 jax_num_cpu_devices：初始化前设置有效，初始化后设置报 RuntimeError
probe = r"""
import jax
jax.config.update('jax_num_cpu_devices', 3)       # 紧跟在 import jax 之后
print(jax.device_count())
try:
    jax.config.update('jax_num_cpu_devices', 5)   # 此时后端已初始化
except RuntimeError as e:
    print('RuntimeError')
"""
env = {k: v for k, v in os.environ.items() if k != "XLA_FLAGS"}   # 去掉本内核的 XLA_FLAGS
out = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, env=env, timeout=120)
print(out.stdout)
assert out.stdout.split() == ["3", "RuntimeError"]
assert jax.config.jax_platforms is None        # 未设置 JAX_PLATFORMS 时为 None（自动选择后端）

# pmap 在 0.12 中是旧接口：docstring 第一句就让你改用 jax.shard_map
first_line = jax.pmap.__doc__.strip().splitlines()[0]
print(first_line)
assert first_line.startswith("Old way of doing parallel map")

# 课程示例：最基本的 pmap SPMD 示例
n_local = jax.local_device_count()
print("可用设备数:", n_local)
num_devices_to_use = 4 if n_local >= 4 else n_local
print(f"pmap 将使用 {num_devices_to_use} 个设备。")

per_device_batch_size = 8
feature_size = 16
global_batch_size = num_devices_to_use * per_device_batch_size

# 主轴 = 设备轴：(设备数, 每设备批大小, 特征数)
sharded_data = jnp.arange(global_batch_size * feature_size, dtype=jnp.float32).reshape(
    (num_devices_to_use, per_device_batch_size, feature_size))
print("输入形状:", sharded_data.shape)


def simple_computation(x):
    return x * 2.0 + 1.0            # 缩放再加常量


parallel_computation = jax.pmap(simple_computation)
result = parallel_computation(sharded_data)
print("输出形状:", result.shape, "| result[0, 0, 0] =", result[0, 0, 0])   # 原值 0 → 0*2+1 = 1
print("实际用到的设备:", sorted(d.id for d in result.sharding.device_set))
assert result.shape == sharded_data.shape and float(result[0, 0, 0]) == 1.0
assert len(result.sharding.device_set) == 4      # 主轴长 4 < 8 个设备时，只用前 4 个设备

# ✅ 现代写法：jit + NamedSharding。数据保持全局形状 (32, 16)，分片信息写在 sharding 里，不需要手动 reshape 出设备轴
mesh4 = jax.make_mesh((4,), ("data",), devices=jax.devices()[:4])      # 和原示例一样，只用 4 个设备
flat = sharded_data.reshape(global_batch_size, feature_size)
flat_s = jax.device_put(flat, NamedSharding(mesh4, P("data")))
res_jit = jax.jit(simple_computation)(flat_s)
print("jit 版类型:", jax.typeof(res_jit))

# ✅ 现代写法：jax.shard_map。函数看到的是每个设备的局部块 (8, 16)
res_sm = jax.jit(jax.shard_map(simple_computation, mesh=mesh4, in_specs=P("data"), out_specs=P("data")))(flat_s)

assert np.array_equal(np.asarray(res_jit).reshape(result.shape), np.asarray(result))
assert np.array_equal(np.asarray(res_sm), np.asarray(res_jit))
print("pmap / jit+sharding / shard_map 结果一致 ✓")

views = {}


def spy(name):
    def f(x):
        views[name] = str(jax.typeof(x))      # 追踪时记录类型
        return x * 2.0 + 1.0
    return f


jax.pmap(spy("pmap"))(sharded_data)
jax.jit(jax.shard_map(spy("shard_map"), mesh=mesh4, in_specs=P("data"), out_specs=P("data")))(flat_s)
jax.jit(spy("jit"))(flat_s)
for k, v in views.items():
    print(f"{k:10s} 视角: {v}")
assert views["jit"] == "float32[32@data,16]"
assert views["shard_map"].startswith("float32[8,16]") and "V:data" in views["shard_map"]
assert views["pmap"].startswith("float32[8,16]")

# 课程示例（用 jax.random.key + split）：准备工作
num_devices = jax.local_device_count()
print("Number of devices:", num_devices)


def predict(params, inputs):
    return jnp.dot(inputs, params["w"]) + params["b"]       # 一个简单的线性变换


key = jax.random.key(0)
k_w, k_b, k_x = jax.random.split(key, 3)
input_dim, output_dim = 10, 5
params = {"w": jax.random.normal(k_w, (input_dim, output_dim)),
          "b": jax.random.normal(k_b, (output_dim,))}
global_batch_size = 32 * num_devices
dummy_data = jax.random.normal(k_x, (global_batch_size, input_dim))
print(f"Parameter shapes: w={params['w'].shape}, b={params['b'].shape}")
print("Global data batch shape:", dummy_data.shape)

# 为 pmap 重塑数据：[设备数, 每设备批大小, 特征数]
batch_per_device = global_batch_size // num_devices
sharded_data = dummy_data.reshape((num_devices, batch_per_device, input_dim))
print("Sharded data shape:", sharded_data.shape)

# params 广播（None），数据沿第 0 轴切开（0）
parallel_predict = jax.pmap(predict, in_axes=(None, 0))
sharded_predictions = parallel_predict(params, sharded_data)
sharded_predictions.block_until_ready()
print("Output predictions shape:", sharded_predictions.shape)
assert sharded_predictions.shape == (num_devices, batch_per_device, output_dim)

# ❌ 旧写法（错误）：省略 in_axes，以为参数会「自动广播」——默认 in_axes=0，w 的第 0 维 10 会被当成设备轴
import contextlib, io
try:
    with contextlib.redirect_stdout(io.StringIO()):     # JAX 报错前会往 stdout 打印一行源码位置，这里屏蔽掉
        jax.pmap(predict)(params, sharded_data)
    raised = False
except ValueError as e:
    raised = True
    msg = str(e)
    print("省略 in_axes 报错:", type(e).__name__)
    print("  ", msg[msg.find("implies"):msg.find("(full shape")].strip())
# b 的第 0 维是 5，被当成「5 个设备」的映射轴；数据的第 0 维却是 8，于是无法整除
assert raised and "partitioned 5 times" in msg

# ✅ 推荐写法 / 现代写法：参数用 P() 复制、数据用 P('data') 分片；函数用全局形状，不需要 reshape（pmap 写法见上一单元的 in_axes=(None, 0)）
mesh = jax.make_mesh((N,), ("data",))
params_rep = jax.device_put(params, NamedSharding(mesh, P()))
data_sh = jax.device_put(dummy_data, NamedSharding(mesh, P("data")))
print("参数类型:", jax.typeof(params_rep["w"]), "| 数据类型:", jax.typeof(data_sh))

pred_modern = jax.jit(predict)(params_rep, data_sh)
print("输出类型:", jax.typeof(pred_modern))

# 与 pmap 结果一致，并与 NumPy float64 对拍
assert np.allclose(np.asarray(pred_modern).reshape(sharded_predictions.shape), np.asarray(sharded_predictions), atol=1e-6)
ref = f64(dummy_data) @ f64(params["w"]) + f64(params["b"])
assert np.allclose(f64(pred_modern), ref, atol=1e-5)
print("pmap == jit+sharding == NumPy float64 ✓")

# 课程示例：psum / pmean / pmax。示例假设有 4 个设备，pmap 主轴长 4 时只用前 4 个设备
def sum_across_devices(local_value):
    return lax.psum(local_value, axis_name="batch_axis")      # 'batch_axis' 必须与 pmap 的 axis_name 一致


def average_across_devices(local_value):
    return lax.pmean(local_value, axis_name="batch_axis")


def get_max_value(local_value):
    return lax.pmax(local_value, axis_name="data_split")


values = jnp.arange(4.)                                        # 每个设备一个值：[0, 1, 2, 3]
r_sum = jax.pmap(sum_across_devices, axis_name="batch_axis")(values)
r_mean = jax.pmap(average_across_devices, axis_name="batch_axis")(values)
r_max = jax.pmap(get_max_value, axis_name="data_split")(jnp.array([2., 5., 1., 4.]))
r_text = jax.pmap(sum_across_devices, axis_name="batch_axis")(jnp.array([2., 3., 1., 4.]))   # 正文里的例子
print("psum :", repr(r_sum))      # 每个设备都得到 0+1+2+3 = 6
print("pmean:", r_mean)           # (0+1+2+3)/4 = 1.5
print("pmax :", r_max)
print("psum([2,3,1,4]):", r_text)
assert type(r_sum).__name__ == "ArrayImpl" and isinstance(r_sum, jax.Array)   # 不再是 DeviceArray
assert np.array_equal(r_sum, [6.] * 4) and np.array_equal(r_mean, [1.5] * 4)
assert np.array_equal(r_max, [5.] * 4) and np.array_equal(r_text, [10.] * 4)

# ✅ 现代写法：shard_map 中直接用网格轴名做集合通信；out_specs=P() 表示结果在各设备上相同（invariant）
mesh4c = jax.make_mesh((4,), ("batch_axis",), devices=jax.devices()[:4])
vals4 = jax.device_put(jnp.arange(4.), NamedSharding(mesh4c, P("batch_axis")))


def reduce_all(x):                    # x 是每个设备的局部块，形状 (1,)
    return lax.psum(x, "batch_axis"), lax.pmean(x, "batch_axis"), lax.pmax(x, "batch_axis")


s, m, mx = jax.jit(jax.shard_map(reduce_all, mesh=mesh4c, in_specs=P("batch_axis"), out_specs=P()))(vals4)
assert float(s[0]) == float(r_sum[0]) == 6.0
assert float(m[0]) == float(r_mean[0]) == 1.5
assert float(mx[0]) == 3.0                          # max(0,1,2,3)

# ✅ 现代写法：jit + sharding，连集合通信都不用写，jnp.sum 就是全局和，编译器自动插入 all-reduce
f_sum = jax.jit(jnp.sum)
assert float(f_sum(vals4)) == 6.0
print("jit 版用到的通信:", collectives(f_sum.lower(vals4).compile()))

mesh8 = jax.make_mesh((N,), ("i",))
v8 = jax.device_put(jnp.arange(8.), NamedSharding(mesh8, P("i")))
jaxpr_txt = str(jax.make_jaxpr(jax.shard_map(lambda x: lax.pmean(x, "i"), mesh=mesh8,
                                             in_specs=P("i"), out_specs=P()))(v8))
print(jaxpr_txt)
assert "psum" in jaxpr_txt and "div" in jaxpr_txt and "pmean" not in jaxpr_txt

# 两种写法编译出的通信相同：都只有一个 all-reduce
f_pmean = jax.jit(jax.shard_map(lambda x: lax.pmean(x, "i"), mesh=mesh8, in_specs=P("i"), out_specs=P()))
f_psum = jax.jit(jax.shard_map(lambda x: lax.psum(x, "i") / N, mesh=mesh8, in_specs=P("i"), out_specs=P()))
assert collectives(f_pmean.lower(v8).compile()) == collectives(f_psum.lower(v8).compile()) == ["all-reduce"]
assert np.allclose(f_pmean(v8), f_psum(v8))

X8 = np.arange(64.).reshape(8, 8)
X8s = jax.device_put(jnp.asarray(X8), NamedSharding(mesh8, P("i")))      # 设备 i 持有第 i 行，局部形状 (1, 8)


def smap(f, in_specs=P("i"), out_specs=P("i")):
    return jax.jit(jax.shard_map(f, mesh=mesh8, in_specs=in_specs, out_specs=out_specs))


# all_gather：tiled=False 新增一个长度 N 的轴，局部 (8,1,8)；tiled=True 沿已有轴拼接，局部 (8,8)
ag = smap(lambda x: lax.all_gather(x, "i", tiled=True))(X8s)
assert np.array_equal(np.asarray(ag), np.tile(X8, (8, 1)))           # 每个设备都拿到完整的 X
ag_shape = jax.eval_shape(jax.shard_map(lambda x: lax.all_gather(x, "i"), mesh=mesh8, in_specs=P("i"), out_specs=P("i")), X8s).shape
print("all_gather 非 tiled 全局形状:", ag_shape, "| tiled:", ag.shape)

# psum_scatter：先逐元素求和（列和），再把结果沿第 1 维切成 8 块，设备 i 只留第 i 块
rs = smap(lambda x: lax.psum_scatter(x, "i", scatter_dimension=1, tiled=True))(X8s)
assert np.array_equal(np.asarray(rs).ravel(), X8.sum(0))
print("psum_scatter 全局形状:", rs.shape, "= 列和", X8.sum(0)[:4], "...")

# all_to_all：每个设备把自己的行切成 8 块分别发给 8 个设备 → 由「按行切分」变成「按列切分」
a2a = smap(lambda x: lax.all_to_all(x, "i", split_axis=1, concat_axis=0, tiled=True), out_specs=P(None, "i"))(X8s)
assert np.array_equal(np.asarray(a2a), X8)
print("all_to_all 前:", jax.typeof(X8s), "→ 后:", jax.typeof(a2a), "（数据不变，分片从行换成列）")
assert a2a.addressable_shards[3].data.shape == (8, 1)
assert np.array_equal(np.asarray(a2a.addressable_shards[3].data).ravel(), X8[:, 3])   # 设备 3 持有第 3 列

# ppermute：环形右移，设备 j 发给 j+1
pp = smap(lambda x: lax.ppermute(x, "i", [(j, (j + 1) % 8) for j in range(8)]))(X8s)
assert np.array_equal(np.asarray(pp), np.roll(X8, 1, axis=0))

# 恒等式：reduce-scatter + all-gather == all-reduce
rsag = smap(lambda x: lax.all_gather(lax.psum_scatter(x, "i", scatter_dimension=1, tiled=True), "i", axis=1, tiled=True))(X8s)
ar = smap(lambda x: lax.psum(x, "i"))(X8s)
assert np.array_equal(np.asarray(rsag), np.asarray(ar))
print("reduce-scatter + all-gather == all-reduce ✓")

xs1d = np.arange(64.) ** 1.5
xs1d_s = jax.device_put(jnp.asarray(xs1d), NamedSharding(mesh8, P("i")))


def stencil(x):                                   # x：本设备的 8 个元素
    to_right = [(j, j + 1) for j in range(7)]     # 设备 j 把最后一个元素发给 j+1（设备 0 收不到 → 0）
    to_left = [(j + 1, j) for j in range(7)]      # 设备 j+1 把第一个元素发给 j（设备 7 收不到 → 0）
    left_halo = lax.ppermute(x[-1:], "i", to_right)
    right_halo = lax.ppermute(x[:1], "i", to_left)
    padded = jnp.concatenate([left_halo, x, right_halo])
    return padded[:-2] + padded[1:-1] + padded[2:]


y_st = smap(stencil)(xs1d_s)
ref_st = np.convolve(np.pad(xs1d, 1), np.ones(3), mode="valid")       # NumPy float64 参考
assert np.allclose(f64(y_st), ref_st, rtol=1e-6)
print("halo 交换模板计算 == NumPy 卷积 ✓；用到的通信:", collectives(smap(stencil).lower(xs1d_s).compile()))

types = {}


def probe_vma(x):
    types["local"] = str(jax.typeof(x))
    types["psum"] = str(jax.typeof(lax.psum(x, "i")))
    types["all_gather"] = str(jax.typeof(lax.all_gather(x, "i", tiled=True)))
    return x


smap(probe_vma)(X8s)
for k, v in types.items():
    print(f"{k:11s}: {v}")
assert "V:i" in types["local"] and "V:i" not in types["psum"] and "V:i" in types["all_gather"]

try:
    smap(lambda x: lax.all_gather(x, "i", tiled=True), out_specs=P())(X8s)
    raised = False
except ValueError as e:
    raised = True
    print("all_gather + out_specs=P() 报错:", str(e)[:90], "...")
assert raised

# ❌ 旧写法（错误）：以为 psum 不写 axis_name 也能隐式工作
try:
    jax.pmap(lambda x: lax.psum(x * 2.0))(jnp.arange(8.))
    raised = False
except TypeError as e:
    raised = True
    print("省略 axis_name 报错:", e)
assert raised

# ✅ 推荐写法 / 课程示例：给映射轴起名 'devices'，psum 显式引用它
def scaled_sum(x):
    scaled_x = x * 2.0
    return lax.psum(scaled_x, axis_name="devices")


n_devices = jax.local_device_count()
print(f"使用 {n_devices} 个设备。")
data = jnp.arange(n_devices, dtype=jnp.float32)
parallel_computation = jax.pmap(scaled_sum, axis_name="devices")   # 名字绑定到输入的第 0 轴
result = parallel_computation(data)
print("输入数据:", data)
print("结果（所有设备上相同）:", result)
expected_sum = jnp.sum(data * 2.0)
print("预期总和:", expected_sum)
assert jnp.allclose(result[0], expected_sum)

# 两个例子：[0,1] → 2；[0,1,2,3] → 12
assert np.array_equal(jax.pmap(scaled_sum, axis_name="devices")(jnp.arange(2.)), [2., 2.])
assert np.array_equal(jax.pmap(scaled_sum, axis_name="devices")(jnp.arange(4.)), [12.] * 4)

# ✅ 现代写法：轴名 'devices' 写在网格上；shard_map 中先求块内和再 psum，jit 版直接写全局求和
mesh_dev = jax.make_mesh((N,), ("devices",))
data_s = jax.device_put(data, NamedSharding(mesh_dev, P("devices")))
res_sm = jax.jit(jax.shard_map(lambda x: lax.psum(jnp.sum(x * 2.0), "devices"), mesh=mesh_dev,
                               in_specs=P("devices"), out_specs=P()))(data_s)
res_jit = jax.jit(lambda x: jnp.sum(x * 2.0))(data_s)
assert float(res_sm) == float(res_jit) == float(result[0]) == 56.0
print("pmap / shard_map / jit 的结果都是", float(res_jit))

# 课程示例：模型与更新步
def loss_fn(params, local_batch):
    predictions = params["w"] * local_batch["x"] + params["b"]
    error = predictions - local_batch["y"]
    return jnp.mean(error ** 2)


def compute_gradients(params, local_batch):
    return jax.grad(loss_fn)(params, local_batch)


def parallel_update_step(params, sharded_batch):
    local_grads = compute_gradients(params, sharded_batch)                     # 每个设备局部计算梯度
    avg_grads = lax.pmean(local_grads, axis_name="data_parallel_axis")        # 沿命名轴平均
    learning_rate = 0.01
    return jax.tree.map(lambda p, g: p - learning_rate * g, params, avg_grads)


n_devices = jax.local_device_count()
params_s = {"w": jnp.ones(()), "b": jnp.zeros(())}
xs_b = jnp.arange(n_devices * 4, dtype=jnp.float32).reshape((n_devices, 4, 1))
ys_b = xs_b * 2.0 + 0.5                                                        # 真实 w=2.0, b=0.5
sharded_batch = {"x": xs_b, "y": ys_b}

# ❌ 旧写法（错误）：省略 in_axes（标量参数无法按默认 in_axes=0 映射），并使用已删除的 jax.tree_map
try:
    jax.pmap(parallel_update_step, axis_name="data_parallel_axis")(params_s, sharded_batch)
    err1 = None
except IndexError as e:
    err1 = e
    print("省略 in_axes:", type(e).__name__, "-", e)
try:
    jax.tree_map(lambda x: x, params_s)
    err2 = None
except AttributeError as e:
    err2 = e
    print("jax.tree_map:", type(e).__name__, "-", e)
assert isinstance(err1, IndexError) and isinstance(err2, AttributeError)

# ✅ 推荐写法（仍用 pmap）：params 用 None 广播，数据沿 0 轴映射；用 jax.tree.map 取设备 0 的结果
p_update_step = jax.pmap(parallel_update_step, axis_name="data_parallel_axis", in_axes=(None, 0))
new_params = p_update_step(params_s, sharded_batch)
new_params_0 = jax.tree.map(lambda x: x[0], new_params)
print("原始参数:", params_s)
print("更新后的参数（设备 0）:", new_params_0)
# pmean 之后各设备的参数完全相同
assert all(np.all(np.asarray(v) == np.asarray(v)[0]) for v in jax.tree.leaves(new_params))

# ✅ 现代写法：jit + NamedSharding。函数里不出现任何轴名，用全局批次的损失，编译器自动插入梯度 all-reduce
mesh_dp = jax.make_mesh((N,), ("data_parallel_axis",))
batch_global = jax.device_put({"x": xs_b.reshape(-1, 1), "y": ys_b.reshape(-1, 1)},
                              NamedSharding(mesh_dp, P("data_parallel_axis")))
params_rep_s = jax.device_put(params_s, NamedSharding(mesh_dp, P()))


@jax.jit
def update_step_global(params, batch):
    grads = jax.grad(loss_fn)(params, batch)             # 全局平均损失的梯度 == 各设备本地梯度的 pmean
    return jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)


new_params_jit = update_step_global(params_rep_s, batch_global)
print("通信:", collectives(update_step_global.lower(params_rep_s, batch_global).compile()))

# NumPy float64 对拍：解析梯度 dL/dw = 2·mean(err·x)，dL/db = 2·mean(err)
xx, yy = f64(xs_b).ravel(), f64(ys_b).ravel()
err = 1.0 * xx + 0.0 - yy
w_ref, b_ref = 1.0 - 0.01 * 2 * np.mean(err * xx), 0.0 - 0.01 * 2 * np.mean(err)
for k, ref_v in [("w", w_ref), ("b", b_ref)]:
    assert np.allclose(float(new_params_0[k]), ref_v, rtol=1e-6)
    assert np.allclose(float(new_params_jit[k]), ref_v, rtol=1e-6)
print(f"pmap == jit+sharding == NumPy：w={w_ref:.6f}, b={b_ref:.6f} ✓")

# vmap 也接受 axis_name：在单设备上验证集合通信的语义
vx = jnp.arange(4.)
v_psum = jax.vmap(lambda x: lax.psum(x, "b"), axis_name="b")(vx)
v_perm = jax.vmap(lambda x: lax.ppermute(x, "b", [(0, 1), (1, 2), (2, 3), (3, 0)]), axis_name="b")(vx)
v_idx = jax.vmap(lambda x: lax.axis_index("b"), axis_name="b")(vx)
print("vmap psum:", v_psum, "| ppermute:", v_perm, "| axis_index:", v_idx)
assert np.array_equal(v_psum, [6.] * 4) and np.array_equal(v_perm, [3., 0., 1., 2.]) and np.array_equal(v_idx, [0, 1, 2, 3])

# 与 shard_map 版本对照（Explicit 网格：输入要先 device_put 成匹配 in_specs 的分片）
mesh_b = jax.make_mesh((4,), ("b",), devices=jax.devices()[:4])
sm_perm = jax.jit(jax.shard_map(lambda x: lax.ppermute(x, "b", [(0, 1), (1, 2), (2, 3), (3, 0)]),
                                mesh=mesh_b, in_specs=P("b"), out_specs=P("b")))(
    jax.device_put(vx, NamedSharding(mesh_b, P("b"))))
assert np.array_equal(np.asarray(sm_perm), np.asarray(v_perm))
print("vmap 与 shard_map 的 ppermute 结果一致 ✓")

# ❌ 旧写法（过时/错误）：jax.experimental.maps、jnp.ones(16.)、.devices 不加括号、reshape(2,2)、嵌套 pmap
problems = {}
try:
    import jax.experimental.maps  # noqa: F401
except ModuleNotFoundError as e:
    problems["experimental.maps"] = f"{type(e).__name__}: {e}"
try:
    jnp.ones(16.)
except TypeError as e:
    problems["jnp.ones(16.)"] = f"{type(e).__name__}: {str(e)[:70]}"
try:
    jnp.arange(jax.local_device_count()).reshape(2, 2)
except TypeError as e:
    problems["arange(8).reshape(2,2)"] = f"{type(e).__name__}: {str(e)[:70]}"
arr = jax.device_put(jnp.arange(4).reshape(2, 2), jax.devices()[0])
problems[".devices 不加括号"] = f"得到 {type(arr.devices).__name__}，而不是设备网格"


from functools import partial


@partial(jax.pmap, axis_name="model")       # 原示例写作 @jax.pmap(axis_name='model')；这里用 partial 写法，等价而且更通用
def inner_op(x_model_shard, params_model_shard):
    return lax.psum(x_model_shard * params_model_shard, axis_name="model") * 2


@partial(jax.pmap, axis_name="data")
def outer_op(x_data_shard, params_data_shard):
    return lax.psum(inner_op(x_data_shard, params_data_shard), axis_name="data")


data_simple = jnp.arange(4.).reshape(2, 2)
params_simple = jnp.ones(4).reshape(2, 2) * 0.5          # 已把 ones(4.) 修正为 ones(4)，只看嵌套 pmap 本身
try:
    outer_op(data_simple, params_simple)
except ValueError as e:
    problems["嵌套 pmap"] = f"{type(e).__name__}: {str(e)[:110]}"
for k, v in problems.items():
    print(f"- {k:24s} → {v}")
assert len(problems) == 5

# ✅ 推荐写法 / 现代写法：嵌套 pmap → 2×2 二维网格 + shard_map。同一个函数里同时引用两个轴名
mesh22 = jax.make_mesh((2, 2), ("data", "model"), devices=jax.devices()[:4])
print(mesh22)


def inner(x, p):
    return lax.psum(x * p, "model") * 2          # 沿 'model' 轴（同一行的 2 个设备）求和


def outer(x, p):
    return lax.psum(inner(x, p), "data")         # 沿 'data' 轴（同一列的 2 个设备）求和


put22 = lambda a: jax.device_put(a, NamedSharding(mesh22, P("data", "model")))
f22 = jax.jit(jax.shard_map(outer, mesh=mesh22, in_specs=(P("data", "model"),) * 2, out_specs=P("data", "model")))

out_simple = np.asarray(f22(put22(data_simple), put22(params_simple)))
print("简化示例输出:\n", out_simple)
assert np.array_equal(out_simple, np.full((2, 2), 6.0))           # 不是 12.0

# 容易算出的 12 是怎么来的：把内层结果（在 model 方向已经复制）再沿两个轴一起 psum，同一个值被数了 2 次
# vma 类型检查会拒绝这种重复计数（对 model 方向不变的值再沿 model 求和），要关掉 check_vma 才能复现 12
try:
    jax.jit(jax.shard_map(lambda x, p: lax.psum(inner(x, p), ("data", "model")), mesh=mesh22,
                          in_specs=(P("data", "model"),) * 2, out_specs=P()))(put22(data_simple), put22(params_simple))
    vma_raised = False
except ValueError as e:
    vma_raised = True
    print("check_vma=True 拒绝重复求和:", str(e)[:80], "...")
f_double = jax.jit(jax.shard_map(lambda x, p: lax.psum(inner(x, p), ("data", "model")), mesh=mesh22,
                                 in_specs=(P("data", "model"),) * 2, out_specs=P(), check_vma=False))
twelve = float(np.asarray(f_double(put22(data_simple), put22(params_simple))).ravel()[0])
print("关闭 check_vma、沿两个轴 psum(1,1,5,5) =", twelve, "← 12.0 是重复计数的结果")
assert vma_raised and twelve == 12.0

# 4 维示例：data (2,2,2,2)，每个设备拿 (1,1,2,2) 的块
data4 = jnp.arange(16.).reshape(2, 2, 2, 2)
params4 = jnp.ones(16).reshape(2, 2, 2, 2) * 0.5
out4 = np.asarray(f22(put22(data4), put22(params4)))
print("4 维示例 output[0, 0]:\n", out4[0, 0])
assert np.array_equal(out4[0, 0], np.asarray(data4).sum(axis=(0, 1)))   # 2·0.5·Σ_ij x_ij = Σ_ij x_ij
assert all(np.array_equal(out4[i, j], out4[0, 0]) for i in range(2) for j in range(2))

# ❌ 旧写法（过时）：jax.jit(device=...) 与 jax.experimental.shard_map 都会发出 DeprecationWarning
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    jax.jit(lambda x: x + 1, device=jax.devices()[1])(1.0)
    from jax.experimental.shard_map import shard_map as _old_shard_map  # noqa: F401
msgs = [str(x.message) for x in w if issubclass(x.category, DeprecationWarning)]
for m_ in msgs:
    print("DeprecationWarning:", m_[:110])
assert any("device argument on jit is deprecated" in m_ for m_ in msgs)
assert any("jax.experimental.shard_map is deprecated" in m_ for m_ in msgs)

# ✅ 推荐写法：用 device_put 把输入放到目标设备，计算就会跟着数据走；用 jax.shard_map 代替 experimental 版本
y_dev = jax.jit(lambda x: x + 1)(jax.device_put(1.0, jax.devices()[1]))
assert y_dev.device == jax.devices()[1]
print("计算在", y_dev.device, "上执行")

mesh2d = jax.make_mesh((4, 2), ("data", "model"))
print(mesh2d)
A = jnp.arange(32.).reshape(8, 4)
cases = [(P(), (8, 4)), (P("data"), (2, 4)), (P(None, "model"), (8, 2)), (P("data", "model"), (2, 2)),
         (P("model", "data"), (4, 1)), (P(("data", "model")), (1, 4))]
for spec, local in cases:
    a = jax.device_put(A, NamedSharding(mesh2d, spec))
    print(f"{str(spec):26s} 局部块 {a.addressable_shards[0].data.shape}   类型 {jax.typeof(a)}")
    assert a.addressable_shards[0].data.shape == local
    assert np.array_equal(np.asarray(a), np.asarray(A))       # 分片只改变存放位置，不改变数值

print("P('data', 'model') 的切分方式：")
jax.debug.visualize_array_sharding(jax.device_put(A, NamedSharding(mesh2d, P("data", "model"))))

xs2 = jax.device_put(jnp.ones((8, 4)), NamedSharding(mesh2d, P("data", "model")))
ws2 = jax.device_put(jnp.ones((4, 6)), NamedSharding(mesh2d, P("model", None)))
print("网格轴类型:", mesh2d.axis_types)
assert mesh2d.axis_types == (AxisType.Explicit, AxisType.Explicit)


@jax.jit
def trace_types(x):
    print("  追踪时 x  :", jax.typeof(x))
    print("  追踪时 x.T:", jax.typeof(x.T))          # 转置后分片也跟着交换
    print("  追踪时 sum(0):", jax.typeof(x.sum(0)))  # 被归约掉的 data 维消失
    return x.T


trace_types(xs2)

# 收缩维 (4) 两边都沿 model 切分 → 输出分片有歧义 → ShardingTypeError
try:
    xs2 @ ws2
    raised = False
except Exception as e:
    raised = type(e).__name__ == "ShardingTypeError"
    print("歧义矩阵乘:", type(e).__name__, "-", str(e)[:95], "...")
assert raised

# 用 out_sharding 明确指定：结果按行沿 data 切分（编译器插入沿 model 的 all-reduce）
r_mm = jnp.matmul(xs2, ws2, out_sharding=NamedSharding(mesh2d, P("data", None)))
print("指定 out_sharding 后:", jax.typeof(r_mm))
assert str(jax.typeof(r_mm)) == "float32[8@data,6]" and np.allclose(np.asarray(r_mm), 4.0)

# Auto 网格：类型中不显示分片，由编译器决定，矩阵乘不会报错
mesh2d_auto = jax.make_mesh((4, 2), ("data", "model"), axis_types=(AxisType.Auto,) * 2)
xa2 = jax.device_put(jnp.ones((8, 4)), NamedSharding(mesh2d_auto, P("data", "model")))
wa2 = jax.device_put(jnp.ones((4, 6)), NamedSharding(mesh2d_auto, P("model", None)))
ra = jax.jit(jnp.matmul)(xa2, wa2)
print("Auto 网格: 类型", jax.typeof(ra), "| 编译器选择的分片", ra.sharding.spec)
assert str(jax.typeof(ra)) == "float32[8,6]" and np.allclose(np.asarray(ra), 4.0)

# Manual：shard_map 函数体内部，网格轴变成 Manual
seen = {}


def peek(x):
    seen["types"] = jax.sharding.get_abstract_mesh().axis_types
    return x


jax.jit(jax.shard_map(peek, mesh=mesh2d, in_specs=P("data", "model"), out_specs=P("data", "model")))(xs2)
print("shard_map 内部的轴类型:", seen["types"])
assert seen["types"] == (AxisType.Manual, AxisType.Manual)

# Auto 网格：with_sharding_constraint 是提示，中间值被重排成 P(None, 'model')
@jax.jit
def f_hint(x):
    y = x * 2
    y = lax.with_sharding_constraint(y, NamedSharding(mesh2d_auto, P(None, "model")))
    return y + 1


xa = jax.device_put(jnp.arange(32.).reshape(8, 4), NamedSharding(mesh2d_auto, P("data", "model")))
r_hint = f_hint(xa)
print("Auto: 输入", xa.sharding.spec, "→ 输出", r_hint.sharding.spec)
assert r_hint.sharding.spec == P(None, "model")
assert np.allclose(np.asarray(r_hint), np.arange(32.).reshape(8, 4) * 2 + 1)

# Explicit 网格：分片不一致时 with_sharding_constraint 当作断言失败
xe = jax.device_put(jnp.arange(32.).reshape(8, 4), NamedSharding(mesh2d, P("data", "model")))
try:
    with jax.set_mesh(mesh2d):
        jax.jit(lambda x: lax.with_sharding_constraint(x * 2, P(None, "model")))(xe)
    raised = False
except AssertionError as e:
    raised = True
    print("Explicit: with_sharding_constraint 报 AssertionError:", str(e)[:100], "...")
assert raised

# Explicit 网格：用 reshard 改变分片，类型随之改变
with jax.set_mesh(mesh2d):
    r_re = jax.jit(lambda x: jax.sharding.reshard(x * 2, P(None, "model")))(xe)
print("Explicit: reshard 后类型", jax.typeof(r_re))
assert str(jax.typeof(r_re)) == "float32[8,4@model]"
assert np.allclose(np.asarray(r_re), np.arange(32.).reshape(8, 4) * 2)