第 5 章　用 pmap 跨设备并行

# 必须在 import jax 之前设置！否则需要重启内核
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import io, contextlib, time, warnings
warnings.simplefilter("default")  # 让 DeprecationWarning 等警告可见，便于检查 API 状态

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P, AxisType

print("JAX 版本:", jax.__version__)
print("设备列表:", jax.devices())
assert jax.device_count() == 8, "没有得到 8 个设备：请重启内核后先运行本单元"

def same(a, b, **kw):
    """比较两个数组是否一致：先转成 numpy，避免不同网格上的数组无法直接运算。"""
    return np.allclose(np.asarray(a), np.asarray(b), **kw)

# 课程示例：检查设备
import jax
import jax.numpy as jnp

print(jax.devices())
print("设备数量:", len(jax.devices()))

# 课程示例（旧接口 pmap）：输入前导维度 = 设备数
num_devices = len(jax.devices())

def simple_computation(x):
    return x * 2

pmapped_computation = jax.pmap(simple_computation)
input_data = jnp.arange(num_devices * 3).reshape((num_devices, 3))

try:
    result = pmapped_computation(input_data)
    print("结果:\n", result)
    print("类型:", type(result))
    print("形状:", result.shape)
    print("所在设备数:", len(result.devices()))
except Exception as e:
    print("出错:", e)

# ✅ 现代写法：jit + NamedSharding —— 不用手动 reshape，像写单设备代码一样
mesh = jax.make_mesh((8,), ("data",))
x_sharded = jax.device_put(jnp.arange(8 * 3), NamedSharding(mesh, P("data")))   # (24,) 第 0 维切到 'data'
print("每设备局部形状:", [s.data.shape for s in x_sharded.addressable_shards])

y_jit = jax.jit(simple_computation)(x_sharded)
print("结果分片:", y_jit.sharding.spec)

# ✅ 现代写法：shard_map —— 写法与 pmap 一一对应
y_sm = jax.jit(jax.shard_map(simple_computation, mesh=mesh, in_specs=P("data"), out_specs=P("data")))(x_sharded)

assert same(y_jit, result.ravel()) and same(y_sm, result.ravel())
print("✓ jit+sharding、shard_map 与 pmap 结果一致")

x83 = jnp.arange(24).reshape(8, 3)
shapes = {}
jax.pmap(lambda x: (shapes.__setitem__("pmap", x.shape), x)[1])(x83)
jax.shard_map(lambda x: (shapes.__setitem__("shard_map", x.shape), x)[1],
              mesh=mesh, in_specs=P("data"), out_specs=P("data"))(jax.device_put(x83, NamedSharding(mesh, P("data"))))
print("设备内看到的局部形状:", shapes)
assert shapes == {"pmap": (3,), "shard_map": (1, 3)}

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    out = jax.pmap(lambda x: x * 2)(jnp.arange(8.0))
print("调用 pmap 产生的警告数:", len(caught))
print("\n--- jax.pmap.__doc__ 开头 ---")
print(jax.pmap.__doc__[:640])
print("--- 返回值 ---")
print("是否 jax.Array:", isinstance(out, jax.Array))
print("sharding:", out.sharding)
assert len(caught) == 0 and isinstance(out, jax.Array)
assert "Old way" in jax.pmap.__doc__ and "_internal_pmap_axis_name" in str(out.sharding)

# ❌ 旧写法（过时）：旧资料常用的 jax.experimental.shard_map 自 v0.8.0 起弃用
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    from jax.experimental.shard_map import shard_map as old_shard_map
print("导入时的警告:", [str(w.message) for w in caught])
assert any(issubclass(w.category, DeprecationWarning) for w in caught)

# ✅ 现代写法：直接用 jax.shard_map
assert hasattr(jax, "shard_map")
print(jax.shard_map)

# 课程示例（旧接口 pmap）：in_axes=(0, None) —— data 沿轴 0 映射，scalar 复制到每个设备
import numpy as np

def scale(x, k):
    # 按标量因子 k 缩放 x
    return x * k

data = jnp.arange(12, dtype=jnp.float32).reshape((4, 3))   # 4 个条目，特征大小为 3
scalar = jnp.float32(10.0)

pmapped_scale = jax.pmap(scale, in_axes=(0, None), out_axes=0)
result = pmapped_scale(data, scalar)

print("设备数量:", jax.local_device_count())
print("原始数据形状:", data.shape)
print("pmap 结果形状:", result.shape)
print("pmap 结果:\n", result)
print("实际参与计算的设备数:", len(result.devices()))
assert len(result.devices()) == 4   # 只用了 4 个设备：每个设备一行 (3,)

# ❌ 错误写法：2 个设备上直接对 (4, 3) 做 pmap（用 devices= 模拟"只有 2 个设备"）
two = jax.devices()[:2]
try:
    jax.pmap(scale, in_axes=(0, None), devices=two)(data, scalar)
    raise AssertionError("不应运行成功")
except ValueError as e:
    print("实测报错:", str(e)[:120])

# 同理：映射轴长 16 > 8 个设备也会报错（0.12 的报错信息并不直观）
try:
    jax.pmap(lambda x: x * 2)(jnp.arange(16.0))
    raise AssertionError("不应运行成功")
except ValueError as e:
    print("16 > 8 报错:", str(e)[:120])

# ✅ 修正代码：先 reshape 成 (N, B/N, ...)，每个设备拿到一个 (B/N, ...) 的小批次
def scale_shape(x, k):
    print("  设备内 x 的形状:", x.shape)   # 追踪时打印一次
    return x * k

per_device = data.reshape(2, 2, 3)                                   # (4, 3) -> (2 设备, 2 行, 3)
r_fixed = jax.pmap(scale_shape, in_axes=(0, None), devices=two)(per_device, scalar)
assert r_fixed.shape == (2, 2, 3)
assert same(r_fixed.reshape(4, 3), data * scalar)
print("✓ reshape 后在 2 个设备上运行正确")

# ✅ 现代写法：jit + NamedSharding —— 这才是"自动切成 B/N 块"的语义，不需要 reshape
mesh2 = jax.make_mesh((2,), ("data",), devices=two)
data_sh = jax.device_put(data, NamedSharding(mesh2, P("data", None)))
print("每设备局部形状:", data_sh.addressable_shards[0].data.shape)   # (2, 3)：每个设备拿到 2 行

r_jit = jax.jit(scale)(data_sh, scalar)
assert data_sh.addressable_shards[0].data.shape == (2, 3)
assert same(r_jit, result) and same(r_jit, r_fixed.reshape(4, 3))
print("✓ jit+sharding 与 pmap 结果一致")

x8 = jnp.arange(16.0).reshape(8, 2)

# 旧接口 pmap：先 pmean 再 out_axes=None（正确用法）
ok = jax.pmap(lambda x: lax.pmean(x.sum(), "d"), axis_name="d", out_axes=None)(x8)

# ❌ 忘了聚合 + out_axes=None：pmap 不报错，默默返回一个设备的局部值 0+1=1
bad = jax.pmap(lambda x: x.sum(), out_axes=None)(x8)
print("正确（pmean 后）:", ok, "   忘记聚合:", bad)
assert float(ok) == 15.0                 # 全局：各设备局部和的平均
assert float(bad) == float(x8[0].sum())  # 静默地等于设备 0 的局部值，没有任何报错

# ✅ 现代写法：shard_map 的 out_specs=P() 会静态检查"是否真的复制"，没聚合就直接报错
xs16 = jax.device_put(jnp.arange(16.0), NamedSharding(mesh, P("data")))
ok_sm = jax.shard_map(lambda x: lax.pmean(x.sum(), "data"), mesh=mesh, in_specs=P("data"), out_specs=P())(xs16)
assert same(ok_sm, ok)
try:
    jax.shard_map(lambda x: x.sum(), mesh=mesh, in_specs=P("data"), out_specs=P())(xs16)
    raise AssertionError("应当报错")
except ValueError as e:
    print("shard_map 报错（好事）:", str(e).splitlines()[0][:110], "...")

# 旧接口 pmap：params 整体复制，x 沿轴 1 映射
params = {"w": jnp.ones((3,)), "b": jnp.float32(0.5)}
xb = jnp.arange(3 * 8, dtype=jnp.float32).reshape(3, 8)   # 设备维放在轴 1

def affine(p, x):
    return x * p["w"] + p["b"]

out_p = jax.pmap(affine, in_axes=(None, 1))(params, xb)     # 设备 i 拿到 xb[:, i]，形状 (3,)
print("pmap 输出形状:", out_p.shape)                         # (8, 3)：设备维堆叠到轴 0

# 也可以对字典内部分别指定：w 复制、b 映射
params2 = {"w": jnp.ones((3,)), "b": jnp.arange(8.0)}
out_p2 = jax.pmap(affine, in_axes=({"w": None, "b": 0}, 1))(params2, xb)

# ✅ 现代写法：jit + NamedSharding —— 直接把 xb 的第 1 维切到 'data' 轴，参数复制
xb_sh = jax.device_put(xb, NamedSharding(mesh, P(None, "data")))
out_m = jax.jit(lambda p, x: (x * p["w"][:, None] + p["b"]).T)(params, xb_sh)
assert same(out_m, out_p)

# 字典内部 b 映射：b 也按 'data' 切分
b_sh = jax.device_put(params2["b"], NamedSharding(mesh, P("data")))
out_m2 = jax.jit(lambda w, b, x: x.T * w + b[:, None])(params2["w"], b_sh, xb_sh)
assert same(out_m2, out_p2)
print("✓ 两种 pytree in_axes 的现代写法与 pmap 一致")

# 课程示例：查看设备（输出因硬件而异）
print(jax.devices())
# 单 CPU: [CpuDevice(id=0)]   多 GPU: [CudaDevice(id=0), ...]（旧版 JAX 显示为 GpuDevice）   TPU: [TpuDevice(id=0), ...]

# 课程示例（旧接口 pmap）：用 axis_name 给并行轴命名
def simple_computation(x):
    return x * 2

data = jnp.arange(4 * 5).reshape((4, 5))   # 每个设备一行（共 4 行）

# pmap 隐式创建一个大小为 4 的一维网格；data 的第 0 轴映射到名为 'data_parallel_axis' 的设备轴
parallel_computation = jax.pmap(simple_computation, axis_name="data_parallel_axis", in_axes=0)
result = parallel_computation(data)
print(result.shape)   # (4, 5)
print(result)

# ✅ 现代写法：shard_map —— 轴名直接写在网格上
mesh4 = jax.make_mesh((4,), ("data_parallel_axis",), devices=jax.devices()[:4])
result_sm = jax.shard_map(simple_computation, mesh=mesh4,
                          in_specs=P("data_parallel_axis"), out_specs=P("data_parallel_axis"))(
    jax.device_put(data, NamedSharding(mesh4, P("data_parallel_axis"))))
assert same(result_sm, result)
print("✓ shard_map 与 pmap 一致；网格:", mesh4)

# ❌ 旧写法（过时）：jax.experimental.maps 已在 0.12 中删除
try:
    from jax.experimental.maps import Mesh as OldMesh
    raise AssertionError("不应导入成功")
except ModuleNotFoundError as e:
    print("实测:", e)

# ✅ 现代写法：jax.make_mesh（推荐，会按硬件拓扑排列设备）或 jax.sharding.Mesh（手动排列）
mesh2d = jax.make_mesh((4, 2), ("data", "model"))
mesh2d_manual = Mesh(np.array(jax.devices()).reshape(4, 2), ("data", "model"))
print(mesh2d)
print("设备排列:\n", mesh2d.devices)
assert mesh2d.shape == mesh2d_manual.shape == {"data": 4, "model": 2}

A = jnp.arange(8 * 4, dtype=jnp.float32).reshape(8, 4)
for spec, local in [(P("data", "model"), (2, 2)), (P("data", None), (2, 4)), (P(None, "model"), (8, 2)),
                    (P(), (8, 4)), (P(("data", "model"), None), (1, 4))]:
    a = jax.device_put(A, NamedSharding(mesh2d, spec))
    assert a.addressable_shards[0].data.shape == local
    print(f"{str(spec):28s} 每设备局部形状 {local}   类型 {jax.typeof(a)}")

mesh_e = jax.make_mesh((8,), ("data",))                              # 默认 Explicit
mesh_a = jax.make_mesh((8,), ("data",), axis_types=(AxisType.Auto,))  # 显式指定 Auto
assert mesh_e.axis_types == (AxisType.Explicit,)

f = lambda x: lax.psum(x.sum(), "data")
raw = jnp.arange(16.0)        # 没有 device_put，未分片
try:
    jax.shard_map(f, mesh=mesh_e, in_specs=P("data"), out_specs=P())(raw)
    raise AssertionError("应当报错")
except ValueError as e:
    print("Explicit 网格 + 未分片输入 报错:", str(e)[:100], "...")

r1 = jax.shard_map(f, mesh=mesh_e, in_specs=P("data"), out_specs=P())(jax.device_put(raw, NamedSharding(mesh_e, P("data"))))
r2 = jax.shard_map(f, mesh=mesh_a, in_specs=P("data"), out_specs=P())(raw)
assert float(r1) == float(r2) == 120.0
print("✓ 方案 1（先 device_put）与方案 2（Auto 网格）都得到 120")

x2d = jax.device_put(A, NamedSharding(mesh2d, P("data", "model")))
try:
    jax.debug.visualize_array_sharding(x2d)
except Exception as e:
    print("visualize_array_sharding 不可用:", e)

def show_sharding(x):
    """不依赖 rich 的简易版：列出每个设备持有的切片范围。"""
    print(f"数组形状 {x.shape}, spec = {x.sharding.spec}")
    for dev, idx in sorted(x.sharding.devices_indices_map(x.shape).items(), key=lambda kv: kv[0].id):
        rng = ", ".join(f"{s.start or 0}:{s.stop if s.stop is not None else n}" for s, n in zip(idx, x.shape))
        print(f"  {dev}: [{rng}]")

show_sharding(x2d)

# ❌ 旧写法（过时）：嵌套 pmap 在 0.12 中报错
xs = jnp.ones((2, 4))
try:
    jax.pmap(jax.pmap(lambda x: lax.psum(x, "j"), axis_name="j"), axis_name="i")(xs)
    raise AssertionError("不应运行成功")
except ValueError as e:
    print("实测报错:", str(e)[:130])

# ✅ 现代写法：(2, 4) 网格 + shard_map，可以只沿某一个轴，或同时沿多个轴做集体操作
mesh_ij = jax.make_mesh((2, 4), ("i", "j"))
xs_sh = jax.device_put(xs, NamedSharding(mesh_ij, P("i", "j")))
psum_j = jax.shard_map(lambda x: lax.psum(x, "j"), mesh=mesh_ij, in_specs=P("i", "j"), out_specs=P("i", None))
psum_ij = jax.shard_map(lambda x: lax.psum(x, ("i", "j")), mesh=mesh_ij, in_specs=P("i", "j"), out_specs=P())
assert same(psum_j(xs_sh), [[4.0], [4.0]])   # 每行 4 个 1 相加
assert same(psum_ij(xs_sh), [[8.0]])         # 全部 8 个 1 相加
print("✓ 沿 j 求和:", psum_j(xs_sh).ravel(), "  沿 (i, j) 求和:", psum_ij(xs_sh).ravel())

# 课程示例（旧接口 pmap）：lax.psum
num_devices = 4   # 示例按 4 个设备写（本机 8 个，pmap 只用前 4 个）

def calculate_local_sum(x):
    local_value = jnp.sum(x)                                  # 每个设备先算本地数据的和
    total_sum = lax.psum(local_value, axis_name="devices")    # 在所有设备上求和
    return total_sum

data = jnp.arange(num_devices * 3).reshape((num_devices, 3))
# 设备 0: [0 1 2] -> 3；设备 1: [3 4 5] -> 12；设备 2: [6 7 8] -> 21；设备 3: [9 10 11] -> 30

pmapped_calculate_sum = jax.pmap(calculate_local_sum, axis_name="devices")   # axis_name 必须与 psum 一致
result_psum = pmapped_calculate_sum(data)
print(result_psum)
assert same(result_psum, [66] * 4)   # 3 + 12 + 21 + 30 = 66

# ✅ 现代写法：shard_map + psum（以及 jit 自动并行）
mesh4 = jax.make_mesh((4,), ("devices",), devices=jax.devices()[:4])
data_sh = jax.device_put(data, NamedSharding(mesh4, P("devices")))

def local_sum_then_psum(x):          # x: 每个设备上的 (1, 3) 块
    local = jnp.sum(x)
    print("  local 的类型:", jax.typeof(local), "  psum 后:", jax.typeof(lax.psum(local, "devices")))
    return lax.psum(local, "devices")

psum_one = jax.jit(jax.shard_map(local_sum_then_psum, mesh=mesh4, in_specs=P("devices"), out_specs=P()))(data_sh)
# 想和 pmap 一样每个设备一份：输出保留一维，out_specs=P('devices')
psum_each = jax.shard_map(lambda x: lax.psum(jnp.sum(x), "devices")[None],
                          mesh=mesh4, in_specs=P("devices"), out_specs=P("devices"))(data_sh)
psum_jit = jax.jit(jnp.sum)(data_sh)   # jit 自动并行：直接写全局 sum

assert int(psum_one) == int(psum_jit) == 66 and same(psum_each, result_psum)
print("✓ shard_map(out_specs=P()) =", psum_one, "；每设备一份 =", psum_each, "；jit =", psum_jit)

# 课程示例（旧接口 pmap）：lax.pmean
def calculate_local_mean(x):
    local_value = jnp.mean(x)
    global_mean = lax.pmean(local_value, axis_name="devices")
    return global_mean

data = jnp.arange(num_devices * 3, dtype=jnp.float32).reshape((num_devices, 3))
# 本地平均值：1.0, 4.0, 7.0, 10.0
pmapped_calculate_mean = jax.pmap(calculate_local_mean, axis_name="devices")
result_pmean = pmapped_calculate_mean(data)
print(result_pmean)
assert same(result_pmean, [5.5] * 4)   # (1 + 4 + 7 + 10) / 4 = 5.5

# ✅ 现代写法：shard_map + pmean，以及 jit 自动并行
data_sh = jax.device_put(data, NamedSharding(mesh4, P("devices")))
pmean_sm = jax.jit(jax.shard_map(lambda x: lax.pmean(jnp.mean(x), "devices"),
                                 mesh=mesh4, in_specs=P("devices"), out_specs=P()))(data_sh)
pmean_jit = jax.jit(jnp.mean)(data_sh)
assert float(pmean_sm) == float(pmean_jit) == float(result_pmean[0]) == 5.5
print("✓ shard_map pmean =", pmean_sm, "；jit mean =", pmean_jit)

# 旧接口 pmap：pmax / pmin / all_gather / axis_index / ppermute
n = 8
xv = jnp.array([3.0, 7.0, 1.0, 5.0, 2.0, 8.0, 6.0, 4.0])

def collectives(v):
    return {
        "pmax": lax.pmax(v, "i"),
        "pmin": lax.pmin(v, "i"),
        "all_gather": lax.all_gather(v, "i"),        # 每个设备都拿到长度为 8 的完整向量
        "axis_index": lax.axis_index("i"),           # 当前设备编号 0..7
        "axis_size": lax.psum(1, "i"),               # psum(1) = 轴上设备数
        "ppermute": lax.ppermute(v, "i", perm=[(j, (j + 1) % n) for j in range(n)]),   # 设备 j 发给 j+1
    }

out_pmap = jax.pmap(collectives, axis_name="i")(xv)
for k, v in out_pmap.items():
    print(f"{k:10s}", v if v.ndim == 1 else f"形状 {v.shape}，设备 0 上: {v[0]}")

# ✅ 现代写法：shard_map —— 每个设备的局部块是 (1,)，输出用 out_specs=P('i') 拼回 (8,) 以便与 pmap 对照
mesh_i = jax.make_mesh((8,), ("i",))
xv_sh = jax.device_put(xv, NamedSharding(mesh_i, P("i")))

@jax.jit
@jax.shard_map(mesh=mesh_i, in_specs=P("i"), out_specs=P("i"))
def collectives_sm(v):                       # v: (1,)
    return {
        "pmax": lax.pmax(v, "i"),
        "pmin": lax.pmin(v, "i"),
        "all_gather": lax.all_gather(v, "i", tiled=True)[None, :],   # (1, 8)，拼回后为 (8, 8)
        "axis_index": lax.axis_index("i")[None],
        "axis_size": jnp.full((1,), lax.axis_size("i")),
        "ppermute": lax.ppermute(v, "i", perm=[(j, (j + 1) % n) for j in range(n)]),
    }

out_sm = collectives_sm(xv_sh)
for k in out_pmap:
    assert same(out_sm[k], out_pmap[k]), k
print("✓ 6 种集体操作的 shard_map 写法与 pmap 完全一致")

# jit 自动并行的等价物：全局 max/min 直接写，all_gather 就是"把数组变成复制的"
assert float(jax.jit(jnp.max)(xv_sh)) == 8.0 and float(jax.jit(jnp.min)(xv_sh)) == 1.0
gathered = jax.reshard(xv_sh, NamedSharding(mesh_i, P()))   # 传 P() 需要先 with jax.set_mesh(mesh_i):
assert gathered.sharding.is_fully_replicated and same(gathered, xv)
print("✓ jit: jnp.max/jnp.min 自动跨设备；reshard 到复制分片相当于 all_gather")

# ❌ 错误写法：jax.device_put(data, devices) —— 想把数据分到设备列表上
num_devices = 2
devices = jax.local_devices()[:num_devices]
data = np.arange(8.0).reshape(num_devices, -1)   # (2, 4)
try:
    sharded_data = jax.device_put(data, devices)
    raise AssertionError("不应运行成功")
except ValueError as e:
    print("实测报错:", str(e).splitlines()[0])
# 原因：device_put 的第二个参数要么是单个设备 / Sharding，要么是与数据结构对应的 pytree，传设备列表不会做"分片"

# ✅ 修正代码：直接把 numpy 数组交给 pmap（它自己负责分发），并先预热再计时
def complex_computation(x):
    y = jnp.sin(x) * jnp.cos(x)
    z = jnp.tanh(y) + jnp.sqrt(jnp.abs(x))
    return z * 2.0

pmap_complex_computation = jax.pmap(complex_computation, devices=devices)
jit_then_pmap = jax.pmap(jax.jit(complex_computation), devices=devices)

def bench(fn, x, reps=50):
    fn(x).block_until_ready()                  # 预热：排除首次编译时间
    t0 = time.perf_counter()
    for _ in range(reps):
        fn(x).block_until_ready()
    return (time.perf_counter() - t0) / reps

result_pj = pmap_complex_computation(data)
print(f"pmap        : {bench(pmap_complex_computation, data) * 1e6:.1f} µs/次")
print(f"pmap(jit(f)): {bench(jit_then_pmap, data) * 1e6:.1f} µs/次（外层 pmap 已负责编译，显式 jit 不带来收益；微秒级差异来自调度开销和测量噪声）")
assert same(result_pj, jit_then_pmap(data))

# ✅ 现代写法：jit + NamedSharding
mesh_dev2 = jax.make_mesh((2,), ("devices",), devices=devices)
data_sh = jax.device_put(data, NamedSharding(mesh_dev2, P("devices")))   # 正确的"分片版 device_put"
result_jit = jax.jit(complex_computation)(data_sh)
assert same(result_jit, result_pj, atol=1e-6)
print("✓ jit+sharding 与 pmap 一致；输出分片:", result_jit.sharding.spec)

# 课程示例（旧接口 pmap）：pmap + value_and_grad + pmean —— 一次数据并行的梯度计算
def predict(params, x):
    w, b = params
    return w * x + b                     # 线性模型 y = w x + b

def loss_fn(params, x_batch, y_batch):
    predictions = predict(params, x_batch)
    error = predictions - y_batch
    return jnp.mean(error ** 2)          # 均方误差

value_and_grad_fn = jax.value_and_grad(loss_fn)

def parallel_update_step(params, x_shards, y_shards):
    loss, grads = value_and_grad_fn(params, x_shards, y_shards)   # 每个设备算本地梯度
    avg_grads = lax.pmean(grads, axis_name="devices")             # 跨设备平均梯度
    avg_loss = lax.pmean(loss, axis_name="devices")
    return avg_loss, avg_grads

# params 复制（None），x / y 沿轴 0 映射；两个输出都已聚合，所以 out_axes=None
pmap_update_step = jax.pmap(parallel_update_step, axis_name="devices",
                            in_axes=(None, 0, 0), out_axes=(None, None), devices=devices)

params = (jnp.array(2.0), jnp.array(1.0))                 # (w, b)
x_data = np.arange(8.0).reshape(num_devices, -1)          # (2, 4)
# 用 jax.random.key 生成噪声（不用未设种子的 np.random.randn），保证结果可复现
noise = jax.random.normal(jax.random.key(0), x_data.shape) * 0.5
y_data = 3 * x_data + 2 + np.asarray(noise)               # y = 3x + 2 + 噪声

avg_loss, avg_grads = pmap_update_step(params, x_data, y_data)
print(f"设备间的平均损失: {avg_loss:.4f}")
print(f"平均梯度 (dw, db): ({avg_grads[0]:.4f}, {avg_grads[1]:.4f})")

# 验证：与单设备整批计算一致
full_loss, full_grads = value_and_grad_fn(params, x_data.ravel(), y_data.ravel())
assert same(avg_loss, full_loss) and same(avg_grads[0], full_grads[0]) and same(avg_grads[1], full_grads[1])

# ✅ 现代写法 1：jit + NamedSharding —— 直接对分片数据求 value_and_grad，不写任何集体操作
xs_sh = jax.device_put(x_data.ravel(), NamedSharding(mesh_dev2, P("devices")))
ys_sh = jax.device_put(y_data.ravel(), NamedSharding(mesh_dev2, P("devices")))
l1, g1 = jax.jit(value_and_grad_fn)(params, xs_sh, ys_sh)

# ✅ 现代写法 2：shard_map —— 对"全局损失 pmean(本地损失)"求导（不要照搬 pmean(grad)，见下一专题）
@jax.jit
@jax.shard_map(mesh=mesh_dev2, in_specs=(P(), P("devices"), P("devices")), out_specs=(P(), P()))
def update_step_sm(params, x, y):
    return jax.value_and_grad(lambda p: lax.pmean(loss_fn(p, x, y), "devices"))(params)

l2, g2 = update_step_sm(params, xs_sh, ys_sh)
for l, g in [(l1, g1), (l2, g2)]:
    assert same(l, avg_loss) and same(g[0], avg_grads[0]) and same(g[1], avg_grads[1])
print("✓ jit+sharding、shard_map 与 pmap 的损失和梯度完全一致")

# 课程示例（旧接口 pmap）：在 pmap 内部使用 vmap
def process_item(item):
    return jnp.tanh(item) * 2.0

def process_batch_of_items(batch):      # 在每个设备上运行
    vectorized_processor = jax.vmap(process_item)
    return vectorized_processor(batch)

pmap_process_batches = jax.pmap(process_batch_of_items, in_axes=0)

data = np.arange(2 * 3 * 4.0).reshape(num_devices, 3, 4)   # (2, 3, 4)
result_pv = pmap_process_batches(data)
print("每个设备的输入数据形状:", data.shape[1:])   # (3, 4)
print("输出结果形状:", result_pv.shape)            # (2, 3, 4)
# pmap 沿大小为 2 的最外层维分发；每个设备拿到 (3, 4)，内部 vmap 沿大小为 3 的轴向量化

# ✅ 现代写法：jit(vmap(...)) + 分片输入 —— 外层 vmap 替代 pmap 的映射，分片决定放到哪个设备
data_sh = jax.device_put(data, NamedSharding(mesh_dev2, P("devices")))
result_mv = jax.jit(jax.vmap(process_batch_of_items))(data_sh)
assert same(result_mv, result_pv)
print("✓ 一致；每设备局部形状:", result_mv.addressable_shards[0].data.shape)

N_DEV = jax.device_count()
mesh8 = jax.make_mesh((N_DEV,), ("data",))

def loss_w(w, x):
    return jnp.mean((x * w - 1.0) ** 2)

xs = jnp.arange(16.0) / 16
xs_sh = jax.device_put(xs, NamedSharding(mesh8, P("data")))
w0 = jnp.array(0.5)

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

ref = jax.grad(loss_w)(w0, xs)                                  # 单设备全局梯度（正确值）
g_pmap = jax.pmap(lambda w, x: lax.pmean(jax.grad(loss_w)(w, x), "d"),
                  axis_name="d", in_axes=(None, 0))(w0, xs.reshape(N_DEV, -1))[0]
# ❌ 照搬 pmap 的写法：grad 后 pmean
g_naive = smap(lambda w, x: lax.pmean(jax.grad(loss_w)(w, x), "data"))(w0, xs_sh)

print(f"单设备正确值 : {float(ref):.6f}")
print(f"pmap         : {float(g_pmap):.6f}")
print(f"shard_map 照搬: {float(g_naive):.6f}   比值 = {float(g_naive / ref):.4f}")
assert same(g_pmap, ref)
assert np.isclose(float(g_naive / ref), N_DEV)    # 恰好是设备数倍

import re
jaxpr_text = str(jax.make_jaxpr(
    jax.shard_map(lambda w, x: jax.grad(loss_w)(w, x), mesh=mesh8, in_specs=(P(), P("data")), out_specs=P())
)(w0, xs_sh))
jaxpr_text = re.sub(r"\x1b\[[0-9;]*m", "", jaxpr_text)   # 去掉终端颜色控制符
lines = jaxpr_text.splitlines()
for ln in lines[:3] + ["    ..."] + [l for l in lines if "psum" in l]:
    print(ln)
assert "pvary" in jaxpr_text and "psum_invariant" in jaxpr_text

# ✅ 写法 1（推荐）：对"全局损失" pmean(本地损失) 求导 —— 语义最清楚，不依赖自动 psum 的细节
g1 = smap(lambda w, x: jax.grad(lambda w: lax.pmean(loss_w(w, x), "data"))(w))(w0, xs_sh)

# ✅ 写法 2：grad 已自动求和，再除以设备数
g2 = smap(lambda w, x: jax.grad(loss_w)(w, x) / lax.axis_size("data"))(w0, xs_sh)

# ✅ 写法 3：先用 pcast 把参数显式转成 varying，恢复 pmap 的"本地梯度"语义，再 pmean
g3 = smap(lambda w, x: lax.pmean(jax.grad(loss_w)(lax.pcast(w, "data", to="varying"), x), "data"))(w0, xs_sh)

# 附：jit + sharding 根本不存在这个问题
g4 = jax.jit(jax.grad(loss_w))(w0, xs_sh)

for name, g in [("写法 1", g1), ("写法 2", g2), ("写法 3", g3), ("jit+sharding", g4)]:
    assert same(g, ref), name
    print(f"✓ {name:12s} {float(g):.6f}")

# ❌ 错误写法：函数体用了 lax.axis_index，却只写了 `from jax import pmap`，没有导入 lax
course_src = """
import jax
import jax.numpy as jnp
from jax import pmap

num_devices = 4
sharded_data = jnp.arange(4 * 10).reshape((num_devices, 10))
replicated_data = jnp.array(5.0)

def my_func(x, y):
    return x * y + lax.axis_index('batch')

pmapped_func = pmap(my_func, in_axes=(0, None), out_axes=0, axis_name='batch')
output = pmapped_func(sharded_data, replicated_data)
"""
try:
    exec(course_src, {})          # 在干净的命名空间中运行这段代码
    raise AssertionError("不应运行成功")
except NameError as e:
    print("实测报错: NameError -", e)

# ✅ 修正代码：补上 from jax import lax
from jax import pmap, lax

num_devices = 4   # 实际使用时应写 jax.local_device_count()
sharded_data = jnp.arange(4 * 10).reshape((num_devices, 10))
replicated_data = jnp.array(5.0)

def my_func(x, y):
    # x 是分片的（每个设备形状 (10,)），y 是复制的（标量）
    return x * y + lax.axis_index("batch")

pmapped_func = pmap(my_func, in_axes=(0, None), out_axes=0, axis_name="batch")

print("分片数据形状:", sharded_data.shape)      # (4, 10)
assert sharded_data.shape[0] <= jax.local_device_count()
print("复制数据形状:", replicated_data.shape)   # ()
output_a = pmapped_func(sharded_data, replicated_data)
print("输出形状:", output_a.shape)              # (4, 10)
assert output_a.shape == (4, 10)
assert same(output_a, sharded_data * 5.0 + jnp.arange(4)[:, None])   # 第 i 行加上设备编号 i

# ✅ 现代写法：shard_map —— 局部块是 (1, 10)，lax.axis_index 用法相同
mesh_b = jax.make_mesh((4,), ("batch",), devices=jax.devices()[:4])
out_sm = jax.shard_map(my_func, mesh=mesh_b, in_specs=(P("batch"), P()), out_specs=P("batch"))(
    jax.device_put(sharded_data, NamedSharding(mesh_b, P("batch"))), replicated_data)
assert same(out_sm, output_a)
print("✓ shard_map 与 pmap 一致")

# ❌ 错误写法：用 Python if 判断追踪值 lax.axis_index(...) == 0
def func_with_debug_print(x):
    intermediate = x * 2
    if lax.axis_index("data_axis") == 0:
        jax.debug.print("设备 0 中间值: {val}", val=intermediate)
    result = lax.psum(intermediate, axis_name="data_axis")
    jax.debug.print("设备 {id} psum 后结果: {res}", id=lax.axis_index("data_axis"), res=result)
    return result

data = jnp.arange(2 * 3).reshape((2, 3))     # [[0 1 2], [3 4 5]]
try:
    pmap(func_with_debug_print, axis_name="data_axis")(data)
    raise AssertionError("不应运行成功")
except jax.errors.TracerBoolConversionError as e:
    print("实测报错: TracerBoolConversionError -", str(e).splitlines()[0])

# ✅ 修正代码：用 lax.cond 根据设备编号决定是否打印
def func_with_debug_print_fixed(x):
    intermediate = x * 2
    # 把 intermediate 作为操作数显式传给分支（不要靠闭包捕获），在 shard_map 中也能用
    lax.cond(lax.axis_index("data_axis") == 0,
             lambda v: jax.debug.print("设备 0 中间值: {val}", val=v),
             lambda v: None,
             intermediate)
    result = lax.psum(intermediate, axis_name="data_axis")
    jax.debug.print("设备 {id} psum 后结果: {res}", id=lax.axis_index("data_axis"), res=result)
    return result

output_b = pmap(func_with_debug_print_fixed, axis_name="data_axis")(data)
jax.effects_barrier()   # 等待 jax.debug.print 的异步输出全部打印完，再打印返回值
print("返回值:\n", output_b)
# 设备 0 的中间值 [0 2 4] + 设备 1 的中间值 [6 8 10] = [6 10 14]
assert same(output_b, [[6, 10, 14], [6, 10, 14]])
assert same(output_b[0], data.sum(0) * 2)       # 用单设备计算核对预期值

# ✅ 现代写法：shard_map 中 jax.debug.print / lax.cond / lax.axis_index 的用法完全相同
mesh_d = jax.make_mesh((2,), ("data_axis",), devices=jax.devices()[:2])
# 注意外面要包 jax.jit：0.12 中不加 jit、逐个算子即时执行的 shard_map 遇到 lax.cond 会报 NotImplementedError
out_dbg = jax.jit(jax.shard_map(lambda x: func_with_debug_print_fixed(x[0])[None],   # 局部块 (1, 3) -> 取出 (3,)
                                mesh=mesh_d, in_specs=P("data_axis"), out_specs=P("data_axis")))(
    jax.device_put(data, NamedSharding(mesh_d, P("data_axis"))))
assert same(out_dbg, output_b)
print("✓ shard_map 版本结果一致:", out_dbg.tolist())

def show_err(title, fn):
    try:
        fn()
        print(f"[{title}] 没有报错")
        return None
    except Exception as e:
        print(f"[{title}] {type(e).__name__}: {str(e).splitlines()[0][:130]}")
        return e

x8 = jnp.arange(8.0)
errs = [
    show_err("轴名拼写不一致     ", lambda: jax.pmap(lambda x: lax.psum(x, "dev"), axis_name="devices")(x8)),
    show_err("在并行函数外调用 psum", lambda: lax.psum(x8, "i")),
    show_err("映射轴长 16 > 8 设备 ", lambda: jax.pmap(lambda x: x)(jnp.arange(16.0))),
    show_err("标量却写 in_axes=0   ", lambda: jax.pmap(lambda x, y: x * y, in_axes=(0, 0))(x8, 2.0)),
    show_err("Python if 判断追踪值 ", lambda: jax.pmap(lambda x: x if x > 0 else -x)(x8)),
]
assert all(e is not None for e in errs)
assert isinstance(errs[0], NameError) and isinstance(errs[1], NameError)

with jax.disable_jit():
    r = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")(jnp.arange(4.0))
assert same(r, [6.0] * 4)
print("\n✓ disable_jit 下 pmap 仍可运行:", r)

# 设备检查
num_devices = jax.device_count()
print(f"Number of devices available: {num_devices}")
print(f"Available devices: {jax.devices()}")

# 课程示例 1（旧接口 pmap）：简单数据并行
def scale_by_two(x):
    print("Compiling and running scale_by_two...")   # 只在追踪时执行
    return x * 2

data_size_per_device = 4
total_data_size = num_devices * data_size_per_device
global_data = jnp.arange(total_data_size)
sharded_data = global_data.reshape((num_devices, data_size_per_device))
print("Global data shape:", global_data.shape)
print("Sharded data shape:", sharded_data.shape)

parallel_scale = jax.pmap(scale_by_two)
result1 = parallel_scale(sharded_data)
print("\nOutput shape:", result1.shape)
print("Output type:", type(result1))
print("Output content:\n", result1)

# 验证：普通 print 在 8 个设备上只打印 1 次；第二次调用命中缓存，0 次
def scale_by_two_counted(x):
    print("Compiling and running...")
    return x * 2

p = jax.pmap(scale_by_two_counted)
buf1, buf2 = io.StringIO(), io.StringIO()
with contextlib.redirect_stdout(buf1):
    p(sharded_data)
with contextlib.redirect_stdout(buf2):
    p(sharded_data)
print("第一次调用打印次数:", buf1.getvalue().count("Compiling"), "  第二次:", buf2.getvalue().count("Compiling"))
assert buf1.getvalue().count("Compiling") == 1 and buf2.getvalue().count("Compiling") == 0

# ✅ 现代写法：jit + NamedSharding —— 不需要 reshape
mesh_dev = jax.make_mesh((num_devices,), ("devices",))
g_sh = jax.device_put(global_data, NamedSharding(mesh_dev, P("devices")))   # (32,) 每设备 4 个
result1_m = jax.jit(lambda x: x * 2)(g_sh)
assert same(result1_m, result1.ravel())
print("✓ 一致；每设备局部形状:", result1_m.addressable_shards[0].data.shape)

# 课程示例 2（旧接口 pmap）：in_axes=(0, None)
def scale_and_add(x, y):
    # process_index 是进程编号（单机恒为 0），不是设备编号，所以这里写 "in process" 而不是 "on device"
    print(f"Compiling and running scale_and_add in process {jax.process_index()}...")
    return x * 2 + y

scalar_y = jnp.float32(100.0)
parallel_scale_add = jax.pmap(scale_and_add, in_axes=(0, None))
result2 = parallel_scale_add(sharded_data, scalar_y)
print("\nInput x shape:", sharded_data.shape, "  Input y:", scalar_y)
print("Output shape:", result2.shape)
print("Output content:\n", result2)    # 每个元素 = x * 2 + 100