第 1 章　JAX 简介

# 全章通用的导入与环境信息
import sys
import time
import warnings
import subprocess

import numpy as np
import jax
import jax.numpy as jnp

print(f"Python : {sys.version.split()[0]}")
print(f"JAX    : {jax.__version__}")
print(f"NumPy  : {np.__version__}")
print(f"默认后端: {jax.default_backend()}")
print(f"设备列表: {jax.devices()}")

# 自证：jax.shard_map 是当前推荐的并行接口，jax.pmap 的文档已标注为旧接口
print(jax.pmap.__doc__.strip().splitlines()[0])
assert hasattr(jax, "shard_map")
assert "shard_map" in jax.pmap.__doc__

# 查看 JAX 追踪得到的中间表示 jaxpr
def f(x):
    return jnp.sin(x) ** 2 + 1.0

print("f 的 jaxpr：")
print(jax.make_jaxpr(f)(1.0))

print("grad(f) 的 jaxpr（多出了 cos 和乘法，就是链式法则）：")
print(jax.make_jaxpr(jax.grad(f))(1.0))

# 自证：d/dx [sin²x + 1] = 2·sin x·cos x
assert jnp.allclose(jax.grad(f)(1.0), 2 * jnp.sin(1.0) * jnp.cos(1.0))

# 检查可用设备（CPU 始终存在）
print(f"JAX 设备: {jax.devices()}")

# NumPy 数组创建
np_array = np.array([1.0, 2.0, 3.0])
print(f"NumPy 数组: {np_array}, 类型: {type(np_array)}")

# JAX 数组创建：类型统一用 jax.Array 识别（具体实现类是 ArrayImpl）
jnp_array = jnp.array([1.0, 2.0, 3.0])
print(f"JAX 数组: {jnp_array}, 类型: {type(jnp_array)}")

assert isinstance(jnp_array, jax.Array)
assert not isinstance(jnp_array, np.ndarray)   # JAX 数组不是 NumPy 数组

# 基本运算的写法完全一样
a_np = np.arange(6).reshape((2, 3))
b_np = np.array([[1, 1, 1], [2, 2, 2]])
c_np = a_np + b_np * 2

a_jnp = jnp.arange(6).reshape((2, 3))
b_jnp = jnp.array([[1, 1, 1], [2, 2, 2]])
c_jnp = a_jnp + b_jnp * 2  # 与 NumPy 语法相同

print(f"NumPy 结果:\n{c_np}")
print(f"JAX 结果:\n{c_jnp}")

# 比较结果是否接近（比较浮点结果时很有用）
print(f"结果接近: {np.allclose(c_np, c_jnp)}")
assert np.allclose(c_np, c_jnp)

np_array = np.array([1, 2, 3])
np_array[0] = 100  # 在 NumPy 中没问题
print(f"修改后的 NumPy 数组: {np_array}")

jnp_array = jnp.array([1, 2, 3])
try:
    # 在 JAX 中会抛出 TypeError
    jnp_array[0] = 100
except TypeError as e:
    print(f"\n原地修改 JAX 数组出错: {e}")

# JAX 的做法：用 .at[index].set(value) 创建一个带新值的新数组
updated_jnp_array = jnp_array.at[0].set(100)
print(f"\n原始 JAX 数组（未改变）: {jnp_array}")
print(f"更新后的 JAX 数组（新对象）: {updated_jnp_array}")

assert jnp.array_equal(jnp_array, jnp.array([1, 2, 3]))
assert jnp.array_equal(updated_jnp_array, jnp.array([100, 2, 3]))

# 演示异步派发：不等待 vs 等待计算完成
x = jax.random.normal(jax.random.key(0), (2000, 2000))
_ = (x @ x).block_until_ready()  # 预热，排除首次调用开销

t0 = time.perf_counter()
y = x @ x                        # 只是派发，立即返回
t_dispatch = time.perf_counter() - t0
y.block_until_ready()            # 等它算完，避免影响下一次计时

t0 = time.perf_counter()
y = (x @ x).block_until_ready()  # 等待真正算完
t_real = time.perf_counter() - t0

print(f"不等待（仅派发）: {t_dispatch * 1e3:8.3f} ms")
print(f"block_until_ready: {t_real * 1e3:8.3f} ms")
print("结论：测速必须用 .block_until_ready()，否则会严重低估耗时。")
assert t_dispatch < t_real

# 检查默认浮点类型
np_float = np.array([1.0]).dtype
jnp_float = jnp.array([1.0]).dtype
print(f"默认 NumPy 浮点类型: {np_float}")
print(f"默认 JAX 浮点类型: {jnp_float}")
assert np_float == np.float64
assert jnp_float == jnp.float32

# ❌ 旧写法（过时）：默认配置下 dtype=jnp.float64 不会生效，会被截断为 float32
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    try:
        jnp_float64 = jnp.array([1.0, 2.0], dtype=jnp.float64)
        print(f"指定 float64，实际 dtype: {jnp_float64.dtype}")
    except Exception as e:
        print(f"出错: {type(e).__name__}: {e}")
for w in caught:
    print(f"实测警告: {w.message}")

assert jnp_float64.dtype == jnp.float32
assert any("truncated to dtype float32" in str(w.message) for w in caught)

# ✅ 推荐写法：在程序启动时开启 jax_enable_x64
# 该配置要在创建任何数组之前设置；为了不影响本 notebook 的其他单元，这里在子进程里演示
x64_demo = '''
import jax
jax.config.update("jax_enable_x64", True)   # 必须在创建数组之前设置
import jax.numpy as jnp
a = jnp.array([1.0, 2.0], dtype=jnp.float64)
assert a.dtype == jnp.float64
assert jnp.array([1.0]).dtype == jnp.float64 and jnp.arange(3).dtype == jnp.int64
print("开启 x64 后，显式 float64:", a.dtype)
print("开启 x64 后，默认浮点:", jnp.array([1.0]).dtype, "，默认整数:", jnp.arange(3).dtype)
'''
r = subprocess.run([sys.executable, "-c", x64_demo], capture_output=True, text=True)
print(r.stdout or r.stderr)
assert r.returncode == 0 and "float64" in r.stdout

# 对照：中途才开启 x64 —— 不报错，但开启前创建的数组仍是 float32
late_demo = '''
import jax, jax.numpy as jnp
early = jnp.ones(2)                          # 开启 x64 之前创建
jax.config.update("jax_enable_x64", True)    # 故意晚一步开启
late = jnp.ones(2, dtype=jnp.float64)
assert early.dtype == jnp.float32 and late.dtype == jnp.float64
print("中途开启 x64：之前创建的数组", early.dtype, "，之后创建的数组", late.dtype)
'''
r = subprocess.run([sys.executable, "-c", late_demo], capture_output=True, text=True)
print(r.stdout or r.stderr)
assert r.returncode == 0 and "float32" in r.stdout

# 弱类型：Python 标量不会把精度提到 64 位
print("NumPy: int64 数组 + 1.5 ->", (np.array([1, 2, 3]) + 1.5).dtype)
print("JAX  : int32 数组 + 1.5 ->", (jnp.array([1, 2, 3]) + 1.5).dtype)
assert (jnp.array([1, 2, 3]) + 1.5).dtype == jnp.float32

# 整数溢出：NumPy 报错，JAX 静默回绕
try:
    print(np.array([1, 2, 3], dtype=np.int8) + 300)
except OverflowError as e:
    print(f"\nNumPy int8 + 300 -> OverflowError: {e}")
wrapped = jnp.array([1, 2, 3], dtype=jnp.int8) + 300
print("JAX   int8 + 300 ->", wrapped, "（300 被截断为 44，结果是错的！）")
assert jnp.array_equal(wrapped, jnp.array([45, 46, 47], dtype=jnp.int8))

a = jnp.arange(5)
try:
    np.arange(5)[10]
except IndexError as e:
    print(f"NumPy 越界: IndexError: {e}")

print(f"JAX a[10]            -> {a[10]}   （被截断到最后一个元素）")
print(f"JAX a.at[10].set(99) -> {a.at[10].set(99)}   （更新被丢弃）")
print(f"JAX a.at[10].get(mode='fill', fill_value=-1) -> {a.at[10].get(mode='fill', fill_value=-1)}")

assert a[10] == 4
assert jnp.array_equal(a.at[10].set(99), a)
assert a.at[10].get(mode="fill", fill_value=-1) == -1

src = np.array([1.0, 2.0])
j = jnp.asarray(src)
src[0] = 99.0
print(f"修改 NumPy 源数组后，JAX 数组仍是: {j}")
assert j[0] == 1.0

back = np.asarray(jnp.ones(2))
print(f"np.asarray(jax 数组) 是否可写: {back.flags.writeable}")
try:
    back[0] = 5
except ValueError as e:
    print(f"写入报错: {e}")
assert not back.flags.writeable
assert np.array(jnp.ones(2)).flags.writeable   # np.array 得到可写副本

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

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

# grad：对 loss_fn 关于第一个参数 (params) 求导
grad_fn = jax.grad(loss_fn)

# jit：编译梯度函数
jit_grad_fn = jax.jit(grad_fn)

# vmap：假设 predict 处理单个输入，vmap 让它能处理一批输入
# 映射 inputs 参数的第 0 轴，params 在所有样本间共享 (None)
batched_predict = jax.vmap(predict, in_axes=(None, 0))

# 构造一个小数据集：y = 3*x0 - 2*x1 + 0.5 + 噪声
key = jax.random.key(0)
k_x, k_noise = jax.random.split(key)
inputs = jax.random.normal(k_x, (8, 2))                          # 8 个样本，每个 2 维特征
true_w, true_b = jnp.array([3.0, -2.0]), 0.5
targets = inputs @ true_w + true_b + 0.1 * jax.random.normal(k_noise, (8,))

params = {'w': jnp.zeros(2), 'b': jnp.array(0.0)}

print("损失:", loss_fn(params, inputs, targets))

grads = grad_fn(params, inputs, targets)
print("梯度（与 params 结构相同的字典）:", grads)

jit_grads = jit_grad_fn(params, inputs, targets)
print("JIT 梯度:", jit_grads)

# 自证 1：jit 前后结果一致
assert jax.tree.all(jax.tree.map(jnp.allclose, grads, jit_grads))

# 自证 2：与解析梯度一致
resid = predict(params, inputs) - targets
assert jnp.allclose(grads['w'], 2 / len(targets) * inputs.T @ resid, atol=1e-5)
assert jnp.allclose(grads['b'], 2 * resid.mean(), atol=1e-5)

# vmap 版本：单样本 predict 作用在一批输入上
true_params = {'w': true_w, 'b': jnp.array(true_b)}
single = predict(true_params, inputs[0])          # 单个样本 -> 标量
batch = batched_predict(true_params, inputs)      # 一批样本 -> (8,)
print("单样本输出形状:", single.shape, "批输出形状:", batch.shape)

# 自证 3：vmap 结果与逐个样本循环一致
assert batch.shape == (8,)
assert jnp.allclose(batch, jnp.stack([predict(true_params, xi) for xi in inputs]))

@jax.jit
def update(params, inputs, targets, lr=0.1):
    grads = jax.grad(loss_fn)(params, inputs, targets)
    # jax.tree.map 对参数字典中的每一项做同样的更新
    return jax.tree.map(lambda p, g: p - lr * g, params, grads)

p = {'w': jnp.zeros(2), 'b': jnp.array(0.0)}
for step in range(201):
    p = update(p, inputs, targets)
    if step % 50 == 0:
        print(f"step {step:3d}  loss = {loss_fn(p, inputs, targets):.5f}")
print("学到的参数:", p, "\n真实参数: w = [3, -2], b = 0.5")

# 自证：收敛到噪声水平，参数接近真实值
assert loss_fn(p, inputs, targets) < 0.01
assert jnp.allclose(p['w'], true_w, atol=0.15)

counter = {"n": 0}

@jax.jit
def impure(x):
    print("  [Python print] 正在追踪, x =", x)        # 追踪时执行，x 是一个 Tracer
    jax.debug.print("  [debug.print] 运行时 x = {}", x)  # 每次运行都会执行
    counter["n"] += 1                                    # 副作用：只在追踪时发生
    return x * 2

print("第 1 次调用 (形状 (2,))"); impure(jnp.ones(2)).block_until_ready()
print("第 2 次调用 (形状 (2,))"); impure(jnp.ones(2)).block_until_ready()
print("第 3 次调用 (形状 (3,))"); impure(jnp.ones(3)).block_until_ready()
print(f"counter = {counter['n']}（调用了 3 次，但只追踪了 2 次）")
assert counter["n"] == 2

# 对照：“只带形状和 dtype 的抽象值”是 jit 的行为；单独 grad 时追踪器带具体数值
def py_if(x):
    return x * x if x > 0 else -x

assert jax.grad(py_if)(-2.0) == -1.0      # 单独 grad：Python if 正常
try:
    jax.jit(py_if)(-2.0)
    jit_if_ok = True
except jax.errors.TracerBoolConversionError:
    jit_if_ok = False
assert not jit_if_ok                       # jit：抽象值无法转成 bool
print("单独 grad 时 Python if 可用；jit 下报 TracerBoolConversionError")

scale = 2.0

@jax.jit
def scaled(x):
    return x * scale   # 读取全局变量 —— 不纯

print("scale=2 时:", scaled(jnp.array(1.0)))
scale = 100.0
print("把 scale 改成 100 后:", scaled(jnp.array(1.0)), " <- 仍然是 2，用的是追踪时的旧值")
assert scaled(jnp.array(1.0)) == 2.0

# 正确做法：把会变化的量作为参数传入
scaled_ok = jax.jit(lambda x, s: x * s)
assert scaled_ok(jnp.array(1.0), scale) == 100.0

loss_val, grads2 = jax.value_and_grad(loss_fn)(params, inputs, targets)
print("损失:", loss_val)
print("梯度:", grads2)
assert jnp.allclose(loss_val, loss_fn(params, inputs, targets))

# 对多个参数求导：argnums=(0, 1) 同时对 params 和 inputs 求导
g_params, g_inputs = jax.grad(loss_fn, argnums=(0, 1))(params, inputs, targets)
print("对 inputs 的梯度形状:", g_inputs.shape)
assert g_inputs.shape == inputs.shape

# 非标量输出会报错
try:
    jax.grad(predict)(params, inputs)
except TypeError as e:
    print("非标量输出求导报错:", str(e).splitlines()[0])

# 自证：从已安装的 jax 包元数据读出 Python 版本要求和可用的安装扩展
import importlib.metadata as md_

meta = md_.metadata("jax")
extras = meta.get_all("Provides-Extra")
print("jax 版本:", meta["Version"])
print("Requires-Python:", meta["Requires-Python"])
print("可用的安装扩展 (pip install 'jax[xxx]'):", extras)

assert meta["Requires-Python"] == ">=3.12"
assert {"cuda12", "cuda13", "tpu"} <= set(extras)
assert "cuda11_cudnn82" not in extras
assert "jaxlib" in " ".join(meta.get_all("Requires-Dist"))   # jaxlib 会作为依赖自动安装

# ❌ 旧写法（过时）：旧式 PRNGKey；另外，示例输出 10.811179161071777 在 JAX 0.12 上复现不了
COURSE_VALUE = 10.811179161071777
try:
    key_old = jax.random.PRNGKey(0)
    x_old = jax.random.normal(key_old, (10,))
    y_course = float(jnp.dot(x_old, x_old))
    print(f"原示例给出的值: {COURSE_VALUE}")
    print(f"实测输出    : {y_course}")
except Exception as e:
    print(f"出错: {type(e).__name__}: {e}")

assert abs(y_course - COURSE_VALUE) > 1e-3   # 示例里的数值复现不了

# ✅ 推荐写法：安装验证脚本（改用 jax.random.key）
try:
    devices = jax.devices()
    print("JAX 检测到以下设备：")
    for i, device in enumerate(devices):
        print(f"{i}: {device.platform.upper()} ({device.device_kind})")

    # 测试一个简单计算
    key = jax.random.key(0)
    x = jax.random.normal(key, (10,))
    y = jnp.dot(x, x)
    print(f"\n成功执行了一个简单的 JAX 操作。结果：{y}")

    # 检查默认后端
    print(f"\n默认后端：{jax.default_backend()}")

except Exception as e:
    print("JAX 验证过程中发生错误：")
    print(e)
    print("\n请检查安装步骤，特别是 GPU 驱动 / CUDA 版本（如果适用）。")

# 自证：结果与 NumPy 独立计算一致；新旧密钥写法得到的随机数相同
assert np.isclose(float(y), np.dot(np.asarray(x), np.asarray(x)), rtol=1e-5)
assert np.isclose(float(y), y_course)

# 补充演示：随机数的具体数值依赖 RNG 配置（在子进程中切换 jax_threefry_partitionable）
rng_demo = '''
import jax, jax.numpy as jnp
jax.config.update("jax_threefry_partitionable", {flag})
x = jax.random.normal(jax.random.key(0), (10,))
print(float(jnp.dot(x, x)))
'''
values = {}
for flag in (True, False):
    r = subprocess.run([sys.executable, "-c", rng_demo.format(flag=flag)], capture_output=True, text=True)
    assert r.returncode == 0, r.stderr
    values[flag] = float(r.stdout)
    print(f"jax_threefry_partitionable={flag!s:5}: dot(x, x) = {values[flag]}")
print(f"原示例给出的值                : {COURSE_VALUE}")

assert values[True] != values[False]                                  # 配置不同，随机数不同
assert all(abs(v - COURSE_VALUE) > 1e-3 for v in values.values())     # 都对不上示例里的数值

jax.print_environment_info()

# 创建数组
print(f"JAX 版本: {jax.__version__}")

# 从列表创建数组
python_list = [1.0, 2.0, 3.0]
jax_array_from_list = jnp.array(python_list)
print(f"从列表创建的数组: {jax_array_from_list}")
print(f"类型: {type(jax_array_from_list)}")
assert isinstance(jax_array_from_list, jax.Array)

# 从嵌套列表创建二维数组
python_nested_list = [[1, 2], [3, 4]]
jax_2d_array = jnp.array(python_nested_list)
print(f"二维数组:\n{jax_2d_array}")

# 全零数组
zeros_array = jnp.zeros((2, 3))  # 形状 (2 行, 3 列)
print(f"全零数组:\n{zeros_array}")

# 全一数组
ones_array = jnp.ones((3,), dtype=jnp.int32)  # 形状 (3,)，整数类型
print(f"全一数组: {ones_array}")
print(f"全一数组数据类型: {ones_array.dtype}")

# 等差整数序列
range_array = jnp.arange(5)  # 类似 Python 的 range(5) -> [0, 1, 2, 3, 4]
print(f"arange 数组: {range_array}")

# 线性等分
linspace_array = jnp.linspace(0, 1, 5)  # 从 0 到 1（包含端点）的 5 个点
print(f"linspace 数组: {linspace_array}")

assert zeros_array.dtype == jnp.float32 and ones_array.dtype == jnp.int32

from jax import random

# 创建初始密钥（推荐写法），种子通常是整数
key = random.key(0)
print(f"初始密钥: {key.dtype}，底层数据 {random.key_data(key)}")

# 生成正态分布随机数。随机函数“使用”密钥，但不会修改它
normal_random_numbers = random.normal(key, shape=(2, 2))
print(f"正态随机数:\n{normal_random_numbers}")

# 要生成更多随机数，先拆分密钥
key, subkey = random.split(key)  # 得到一个新密钥和一个子密钥
uniform_random_numbers = random.uniform(subkey, shape=(3,))
print(f"\n拆分后的密钥: {random.key_data(key)}")
print(f"子密钥: {random.key_data(subkey)}")
print(f"均匀随机数: {uniform_random_numbers}")

# 用“相同”的原始密钥再调用一次，结果完全相同 —— 体现纯度和可复现性
same_normal_numbers = random.normal(random.key(0), shape=(2, 2))
print(f"\n相同的正态随机数:\n{same_normal_numbers}")
assert jnp.allclose(normal_random_numbers, same_normal_numbers)

# key 与 PRNGKey 对比演示（除 1.4 的 ❌ 对照单元外，PRNGKey 只出现在这里）
old_key = jax.random.PRNGKey(0)
new_key = jax.random.key(0)
print(f"PRNGKey(0): {old_key}, dtype={old_key.dtype}, shape={old_key.shape}")
print(f"key(0)    : dtype={new_key.dtype}, shape={new_key.shape}")
print(f"key_data(key(0)) = {jax.random.key_data(new_key)}")

# 生成的随机数相同
same = jnp.array_equal(random.normal(old_key, (3,)), random.normal(new_key, (3,)))
print("随机数相同:", same)

# 批量拆分后的形状
print("split(PRNGKey, 3).shape =", random.split(old_key, 3).shape)
print("split(key, 3).shape     =", random.split(new_key, 3).shape)

# 旧式密钥能被误当成数据运算，新式密钥会报错
print("旧式 key + 1 =", old_key + 1, "（静默成功，很危险）")
try:
    new_key + 1
except TypeError as e:
    print("新式 key + 1 报错:", str(e).splitlines()[0])

assert same
assert old_key.dtype == jnp.uint32 and old_key.shape == (2,)
assert new_key.shape == () and jax.dtypes.issubdtype(new_key.dtype, jax.dtypes.prng_key)
assert random.split(new_key, 3).shape == (3,)
assert bool(jax.random.wrap_key_data(old_key) == new_key)   # 新旧互转

# 两种写法都能指定 RNG 算法，但只有类型化密钥会把算法记在 dtype 里
raw_rbg = jax.random.PRNGKey(0, impl="rbg")
typed_rbg = jax.random.key(0, impl="rbg")
print("PRNGKey(0, impl='rbg'):", raw_rbg.dtype, raw_rbg.shape, "| key(0, impl='rbg'):", typed_rbg.dtype)
assert raw_rbg.dtype == jnp.uint32 and raw_rbg.shape == (4,)
assert "rbg" in str(typed_rbg.dtype)
try:
    random.normal(raw_rbg, (2,))      # 原始数据不带算法信息：按默认 threefry 解释，形状对不上
    raw_ok = True
except TypeError:
    raw_ok = False
assert not raw_ok
random.normal(typed_rbg, (2,))        # 类型化密钥自带算法信息，直接可用

key = jax.random.key(42)

# 错误：同一个密钥用了两次
w1 = random.normal(key, (3,))
w2 = random.normal(key, (3,))
print("复用密钥 -> 两组“随机”数完全相同:", jnp.array_equal(w1, w2))
assert jnp.array_equal(w1, w2)

# 正确：拆分出独立的子密钥
key, k1, k2 = random.split(key, 3)
w1, w2 = random.normal(k1, (3,)), random.normal(k2, (3,))
print("拆分子密钥 -> 不同:", not jnp.array_equal(w1, w2))
assert not jnp.array_equal(w1, w2)

# 循环中用 fold_in 按步数派生密钥
base = jax.random.key(0)
for step in range(3):
    print(f"step {step}:", random.uniform(jax.random.fold_in(base, step), (2,)))

# 换成不同的分布也不行：同一密钥的 uniform 与 normal 排序完全一致（单调对应）
u_same, z_same = random.uniform(key, (1000,)), random.normal(key, (1000,))
print("同一密钥的 uniform 与 normal 排序一致:", bool(jnp.all(jnp.argsort(u_same) == jnp.argsort(z_same))))
assert jnp.all(jnp.argsort(u_same) == jnp.argsort(z_same))

x = jnp.arange(12).reshape((3, 4))
print(f"数组 x:\n{x}")

print(f"形状: {x.shape}")      # 各维度大小组成的元组 (行数, 列数)
print(f"数据类型: {x.dtype}")  # 元素类型，例如 float32、int32
print(f"维度数量: {x.ndim}")   # 轴的数量（矩阵为 2）
print(f"元素总数: {x.size}")   # 元素总数 3 * 4 = 12

assert x.shape == (3, 4) and x.dtype == jnp.int32 and x.ndim == 2 and x.size == 12

# ❌ 旧写法（过时）：代码用到 np 却没有 import numpy as np，直接运行报 NameError
course_src = '''
numpy_array = np.array([1, 2, 3])
numpy_array[0] = 100
'''
try:
    exec(course_src, {})   # 在空白命名空间中执行，模拟把这段代码单独复制出去运行
    course_ok = True
except NameError as e:
    course_ok = False
    print(f"实测: NameError: {e}")
assert not course_ok

# ✅ 推荐写法：先导入 NumPy（本教程开头已经统一 import numpy as np）
import numpy as np

numpy_array = np.array([1, 2, 3])
numpy_array[0] = 100  # NumPy 中正常
print(f"修改后的 NumPy 数组: {numpy_array}")

jax_array = jnp.array([1, 2, 3])
try:
    # 会报错，因为 JAX 数组不可变
    jax_array[0] = 100
except TypeError as e:
    print(f"\n尝试原地修改 JAX 数组时出错: {e}")

# 改用函数式的“索引更新”语法，返回新数组
updated_jax_array = jax_array.at[0].set(100)
print(f"\n原始 JAX 数组（未改变）: {jax_array}")
print(f"更新后的 JAX 数组（新对象）: {updated_jax_array}")

# 还可以做更多种类的更新
incremented_array = jax_array.at[1].add(10)  # 索引 1 处的元素加 10
print(f"递增后的数组: {incremented_array}")

assert jnp.array_equal(jax_array, jnp.array([1, 2, 3]))
assert jnp.array_equal(updated_jax_array, jnp.array([100, 2, 3]))
assert jnp.array_equal(incremented_array, jnp.array([1, 12, 3]))

x = jnp.zeros(5)
print("set     :", x.at[1:3].set(7.0))
print("multiply:", jnp.arange(5.).at[::2].multiply(10))
print("max     :", jnp.arange(5.).at[:].max(2.0))
print("apply   :", jnp.arange(5.).at[2].apply(jnp.sqrt))

# 重复下标：JAX 累加，NumPy 不累加
idx = jnp.array([0, 0, 2])
jax_acc = jnp.zeros(3).at[idx].add(1)
print("\nJAX  .at[[0,0,2]].add(1):", jax_acc)
a_np = np.zeros(3); a_np[[0, 0, 2]] += 1
print("NumPy a[[0,0,2]] += 1   :", a_np)
b_np = np.zeros(3); np.add.at(b_np, [0, 0, 2], 1)
print("NumPy np.add.at          :", b_np)
assert jnp.array_equal(jax_acc, jnp.array([2., 0., 1.]))
assert np.array_equal(a_np, [1., 0., 1.]) and np.array_equal(b_np, jax_acc)

# jit 内用 jnp.where 代替布尔掩码赋值
relu = jax.jit(lambda v: jnp.where(v > 0, v, 0.0))
out = relu(jnp.array([-1.0, 2.0, -3.0, 4.0]))
print("\njnp.where 实现 ReLU:", out)
assert jnp.array_equal(out, jnp.array([0., 2., 0., 4.]))

# 原地更新：检查编译后的 HLO 里有没有 copy 指令
def hlo_has_copy(fn, *args):
    return " copy(" in fn.lower(*args).compile().as_text()

v0 = jnp.zeros(1000)
assert not hlo_has_copy(jax.jit(lambda v: (v * 2).at[0].set(1.0)), v0)             # 更新中间结果：原地
assert hlo_has_copy(jax.jit(lambda v: v.at[0].set(1.0)), v0)                        # 更新输入参数：先拷贝
assert not hlo_has_copy(jax.jit(lambda v: v.at[0].set(1.0), donate_argnums=0), v0)  # 捐赠输入后：原地
print("\nHLO 检查：更新中间结果无 copy；更新输入参数有 copy；donate_argnums 后无 copy")

a = jnp.array([[1, 2], [3, 4]])
b = jnp.array([[5, 6], [7, 8]])

# 元素级运算
print(f"a + b:\n{a + b}")
print(f"a * b:\n{a * b}")  # 元素级乘法

# 矩阵乘法
print(f"矩阵积 (jnp.dot):\n{jnp.dot(a, b)}")
print(f"矩阵积 (@):\n{a @ b}")

# 通用函数
print(f"a 的正弦:\n{jnp.sin(a)}")
print(f"a 的指数:\n{jnp.exp(a)}")

assert jnp.array_equal(a @ b, jnp.array([[19, 22], [43, 50]]))
assert np.allclose(jnp.exp(a), np.exp(np.asarray(a)))

data = jnp.arange(10)
print(f"原始数据: {data}")

# 获取单个元素
print(f"索引 3 处的元素: {data[3]}")

# 获取切片
print(f"从索引 2 到 5 的切片: {data[2:5]}")

# 多维索引
matrix = jnp.arange(9).reshape((3, 3))
print(f"矩阵:\n{matrix}")
print(f"第 1 行第 2 列的元素: {matrix[1, 2]}")
print(f"前两行:\n{matrix[:2, :]}")
print(f"第一列:\n{matrix[:, 0]}")

assert data[3] == 3 and matrix[1, 2] == 5
assert jnp.array_equal(matrix[:, 0], jnp.array([0, 3, 6]))

# ❌ 旧写法（过时）：x.device() 已改为属性；用 except AttributeError 兜底也接不住新的错误
x = jnp.ones(3)
escaped = None
try:
    try:
        print(f"数组 x 的设备类型: {x.device().platform}")
    except AttributeError:
        print("设备信息可能需要特定上下文或检查 jax.devices()。")   # 兜底分支，不会被执行
except TypeError as e:
    escaped = e
    print(f"TypeError 没有被 except AttributeError 接住：{e}")
assert isinstance(escaped, TypeError)

# ✅ 推荐写法
x = jnp.ones(3)
print(f"x.device           -> {x.device}（属性，不加括号）")
print(f"x.device.platform  -> {x.device.platform}")
print(f"x.devices()        -> {x.devices()}（方法，返回设备集合）")
print(f"可用设备: {jax.devices()}")

assert isinstance(x.device, jax.Device)
assert x.devices() == {x.device}
assert x.device.platform == jax.default_backend()

# 列出所有可用设备
devices = jax.devices()
print(f"可用设备: {devices}")

# 默认后端
print(f"默认后端: {jax.default_backend()}")

# 每个设备的详细信息
for d in devices:
    print(f"  id={d.id}, platform={d.platform}, kind={d.device_kind}")
assert all(d.platform == jax.default_backend() for d in devices)

x = jnp.arange(10.0)
print(f"数组 x 位于: {x.device}")
assert x.device == jax.devices()[0]

# ❌ 旧写法（过时）：Array.device 已从方法改为属性，加括号调用会报 TypeError
random_array = jax.random.uniform(jax.random.key(42), shape=(2, 2))
try:
    print("设备:", random_array.device())
    old_api_works = True
except TypeError as e:
    old_api_works = False
    print(f"实测: TypeError: {e}")
assert not old_api_works

# ✅ 推荐写法：三种查看方式
print("random_array.device    ->", random_array.device)
print("random_array.devices() ->", random_array.devices())
print("random_array.sharding  ->", random_array.sharding)

assert random_array.devices() == {random_array.device}
assert random_array.sharding.device_set == random_array.devices()
assert random_array.sharding.is_fully_replicated   # 单设备数组：完整数据都在这一个设备上

devices = jax.devices()
x = jnp.arange(10.0)

# 放到第一个设备
x_dev0 = jax.device_put(x, devices[0])
print(f"x_dev0 位于: {x_dev0.device}")

# 放到第二个设备（需要至少两个设备）
try:
    x_dev1 = jax.device_put(x, devices[1])
    print(f"x_dev1 位于: {x_dev1.device}")
except IndexError:
    print(f"只有 {len(devices)} 个设备，无法放到 devices[1]")

# 显式放到 CPU（有 GPU 时很有用）
cpu_device = jax.devices("cpu")[0]
x_cpu = jax.device_put(x, cpu_device)
print(f"x_cpu 位于: {x_cpu.device}")
assert x_cpu.device == cpu_device

# 不指定设备时，新数组在默认设备上
y = jnp.ones(5)
print(f"y 位于默认设备: {y.device}")

# 查询不存在的后端会抛 RuntimeError
try:
    jax.devices("gpu")
except RuntimeError as e:
    print(f"\njax.devices('gpu') 报错: {e}")

multi_dev_demo = '''
import jax
jax.config.update("jax_num_cpu_devices", 4)   # 模拟 4 个 CPU 设备，必须在使用 JAX 之前设置
import jax.numpy as jnp

devs = jax.devices()
assert len(devs) == 4
print("模拟的设备:", devs)

x = jnp.arange(4.0)
a = jax.device_put(x, devs[1])        # committed 到 cpu:1
b = jax.device_put(x, devs[2])        # committed 到 cpu:2
print("a 在", a.device, "，b 在", b.device)

assert (a + 1).device == devs[1]
assert (a + jnp.ones(4)).device == devs[1]          # uncommitted 数组被自动搬到 cpu:1
print("a + jnp.ones(4)（uncommitted 数组被自动搬过来）在", (a + jnp.ones(4)).device)

raised = False
try:
    a + b
except ValueError as e:
    raised = True
    print("a + b 报错 ValueError:", str(e)[:70], "...")
assert raised

# 正确做法：先把数据显式搬到同一设备
c = a + jax.device_put(b, devs[1])
assert c.device == devs[1] and jnp.array_equal(c, 2 * x)
print("搬到同一设备后 a + b =", c, "，位于", c.device)
print("ALL_OK")
'''
r = subprocess.run([sys.executable, "-c", multi_dev_demo], capture_output=True, text=True)
print(r.stdout or r.stderr)
assert r.returncode == 0 and "ALL_OK" in r.stdout

x = jnp.arange(4.0) * 2

host_np = np.asarray(x)
host_get = jax.device_get(x)
print("np.asarray     ->", type(host_np), host_np)
print("jax.device_get ->", type(host_get), host_get)
print("标量 .item()   ->", x.sum().item(), type(x.sum().item()))
assert isinstance(host_get, np.ndarray) and isinstance(x.sum().item(), float)

with jax.default_device(jax.devices("cpu")[0]):
    z = jnp.zeros(2)
print("default_device 块内创建的数组位于:", z.device)
assert z.device == jax.devices("cpu")[0]

# 综合练习：数组的创建、属性与设备
# numpy 通常用于对比，或用于 JAX 不支持的操作（已在开头导入）
print(f"JAX 版本: {jax.__version__}")
print(f"JAX 默认后端: {jax.default_backend()}")
# 检查可用设备（取决于你的环境）
print(f"可用 JAX 设备: {jax.devices()}")

# 从 Python 列表创建
py_list = [1.0, 2.5, 3.0, 4.2]
jax_array_from_list = jnp.array(py_list)
print("从列表创建的 JAX 数组:", jax_array_from_list)
print("类型:", type(jax_array_from_list))
print("数据类型:", jax_array_from_list.dtype)

# 自证：类型是 jax.Array，CPU 上默认也是 float32
assert isinstance(jax_array_from_list, jax.Array)
assert "DeviceArray" not in type(jax_array_from_list).__name__
assert jax_array_from_list.dtype == jnp.float32

# 创建特定数组
zeros_array = jnp.zeros((2, 3))  # 形状 (2 行, 3 列)
print("\n零数组:\n", zeros_array)

ones_array = jnp.ones((3, 2), dtype=jnp.int32)  # 指定数据类型
print("\n全一数组 (int32):\n", ones_array)
print("数据类型:", ones_array.dtype)

range_array = jnp.arange(0, 10, 2)  # 起始值, 停止值（不含）, 步长
print("\n范围数组:", range_array)

linspace_array = jnp.linspace(0, 1, 5)  # 起始值, 停止值（含）, 点数
print("\n等差数组:", linspace_array)

assert jnp.array_equal(range_array, jnp.array([0, 2, 4, 6, 8]))
assert jnp.allclose(linspace_array, jnp.array([0, .25, .5, .75, 1.]))

# 创建 PRNG 密钥
key = jax.random.key(42)  # 种子为 42

# 生成 [0, 1) 均匀分布随机数
random_array = jax.random.uniform(key, shape=(2, 2))
print("随机数组 (均匀分布):\n", random_array)

# 重要提示：要获得新的随机数，必须先“拆分”密钥
key, subkey = jax.random.split(key)
random_normal_array = jax.random.normal(subkey, shape=(3,))  # 标准正态分布
print("\n随机数组 (正态分布):\n", random_normal_array)

assert ((random_array >= 0) & (random_array < 1)).all()

print("检查 random_array:")
print("形状:", random_array.shape)
print("大小:", random_array.size)   # 元素总数
print("维度数:", random_array.ndim)
print("数据类型:", random_array.dtype)

# 检查数组所在的设备（属性，不加括号；旧写法 .device() 的对照见 1.6）
print("设备:", random_array.device)

assert random_array.shape == (2, 2) and random_array.size == 4 and random_array.ndim == 2
assert random_array.dtype == jnp.float32

a = jnp.array([1, 2, 3])
b = jnp.array([4, 5, 6])

# 按元素运算
c_add = a + b
print("a + b =", c_add)

c_mul = a * b
print("a * b =", c_mul)

# 与标量运算（广播）
c_scalar_add = a + 10
print("a + 10 =", c_scalar_add)

c_scalar_mul = a * 2
print("a * 2 =", c_scalar_mul)

# 检查 a 是否保持不变（不可变性）
print("原始 'a':", a)
assert jnp.array_equal(a, jnp.array([1, 2, 3]))
assert jnp.array_equal(c_mul, jnp.array([4, 10, 18]))

# 矩阵乘法
mat_a = jnp.array([[1, 2], [3, 4]])
mat_b = jnp.array([[5, 6], [7, 8]])

mat_prod = mat_a @ mat_b               # 等价于 jnp.matmul(mat_a, mat_b)
print("矩阵乘积 mat_a @ mat_b:\n", mat_prod)
assert jnp.array_equal(mat_prod, jnp.matmul(mat_a, mat_b))
assert np.array_equal(mat_prod, np.asarray(mat_a) @ np.asarray(mat_b))

data = jnp.arange(12).reshape((3, 4))
print("原始数据:\n", data)

# 单个元素
print("\n第 1 行第 2 列的元素:", data[1, 2])

# 整行
print("第 0 行:", data[0, :])

# 整列
print("第 1 列:", data[:, 1])

# 子数组
print("子数组（前 2 行，第 1~2 列）:\n", data[:2, 1:3])

# 直接赋值不被支持
try:
    data[0, 0] = 99
except TypeError as e:
    print("\n直接赋值报错:", str(e).splitlines()[0])

# 函数式更新：返回新数组，这种写法才能和 jit、grad 等变换兼容
updated_data = data.at[0, 0].set(99)
print("\n更新后的数据:\n", updated_data)
print("原始数据（未变）:\n", data)

assert data[0, 0] == 0 and updated_data[0, 0] == 99
assert jnp.array_equal(data[:2, 1:3], jnp.array([[1, 2], [5, 6]]))