第 5 章　互操作性与自定义操作

# 公共准备：导入、版本检查、C++ 目录、计时与对齐分配工具
import sys, time, ctypes, subprocess, warnings, shutil
from pathlib import Path
from functools import partial
import numpy as np
import jax
import jax.numpy as jnp

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| 设备", jax.devices())
assert jax.__version__.startswith("0.12")
assert jax.default_backend() == "cpu", "本章的示例按 CPU 编写（在 Colab 上请使用 CPU 运行时）"

NB_DIR = Path.cwd()                 # notebook 所在的目录
CPP_DIR = NB_DIR / "cpp"            # C++ 源码与 .so 输出目录
CPP_DIR.mkdir(exist_ok=True)
def short_path(s):
    """打印时把本机绝对路径缩短（notebook 目录 → 相对路径，虚拟环境 → <venv>），输出里不出现个人目录"""
    return str(s).replace(str(NB_DIR) + "/", "").replace(sys.prefix, "<venv>")

print("C++ 目录：", short_path(CPP_DIR))

def bench(f, *args, n=20):
    """平均耗时（毫秒）：先预热一次，再循环 n 次，最后 block_until_ready"""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n * 1e3

def aligned_zeros(n, dtype=np.float32, align=64):
    """分配起始地址按 align 字节对齐的一维 NumPy 数组（DLPack 零拷贝导入 JAX 的前提）"""
    itemsize = np.dtype(dtype).itemsize
    buf = np.zeros(n * itemsize + align, dtype=np.uint8)
    off = (-buf.ctypes.data) % align
    out = buf[off:off + n * itemsize].view(dtype)
    assert out.ctypes.data % align == 0
    return out

def shares_memory(j, n):
    """JAX 数组 j 与 NumPy 数组 n 是否指向同一块内存（比较起始指针）"""
    return j.unsafe_buffer_pointer() == n.ctypes.data

# ✅ 推荐写法：NumPy → JAX（对应课程代码块 1；设备用 .device 属性）
numpy_arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
print("原始 NumPy 类型:", type(numpy_arr), "（隐含在主机 CPU 上）")

jax_arr = jnp.array(numpy_arr)
print("转换后类型:", type(jax_arr), "| isinstance(jax.Array):", isinstance(jax_arr, jax.Array))
print("所在设备:", jax_arr.device)

assert isinstance(jax_arr, jax.Array) and jax_arr.dtype == jnp.float32
assert jax_arr.device == jax.devices()[0]

# ❌ 旧写法（过时/错误）：jax_arr.device() —— 0.12 中 .device 是属性（Device 对象），不能调用
try:
    print(jax_arr.device())
except TypeError as e:
    print("TypeError:", e)

# ✅ 推荐写法：.device 属性（单设备数组）；.devices() 返回设备集合（分片数组也适用）
print(jax_arr.device, jax_arr.devices())
assert jax_arr.devices() == {jax_arr.device}

# 课程代码块 2：用 jax.device_put 显式放置
available_devices = jax.devices()
print("可用设备:", available_devices)
target_device = available_devices[1] if len(available_devices) > 1 else available_devices[0]
jax_arr_explicit = jax.device_put(numpy_arr, device=target_device)
print("显式放置后的设备:", jax_arr_explicit.device)
assert jax_arr_explicit.device == target_device

# ❌ 旧写法（过时/错误）：jax.random.PRNGKey(0) 返回 uint32[2] 的“旧式原始 key”，0.12 推荐类型化 key
try:
    old_key = jax.random.PRNGKey(0)
    print("旧式 key:", old_key, "| dtype:", old_key.dtype, "| shape:", old_key.shape)
except Exception as e:
    print(type(e).__name__, e)

# ✅ 推荐写法：jax.random.key 生成类型化 key；JAX → NumPy（课程代码块 3）
key = jax.random.key(0)
print("类型化 key:", key, "| dtype:", key.dtype)
assert jax.dtypes.issubdtype(key.dtype, jax.dtypes.prng_key)

jax_arr = jax.random.normal(key, (3,)) * 2.0 + 1.0
print("JAX 数组:", jax_arr, "| 设备:", jax_arr.device)
numpy_result = np.array(jax_arr)     # 同步点：等待计算完成，并复制回主机
print("NumPy 类型:", type(numpy_result), "| 值:", numpy_result)
assert isinstance(numpy_result, np.ndarray)
np.testing.assert_array_equal(numpy_result, jax_arr)

# 动手验证：各种转换到底复制不复制（用指针比较判断）
big = aligned_zeros(1 << 20); big[:] = 1.0           # 4 MB，64 字节对齐
a_copy = jnp.array(big)
a_asarray = jnp.asarray(big)
a_put = jax.device_put(big)
a_put_noalias = jax.device_put(big, may_alias=False)
for name, arr in [("jnp.array", a_copy), ("jnp.asarray", a_asarray),
                  ("jax.device_put", a_put), ("device_put(may_alias=False)", a_put_noalias)]:
    print(f"{name:28s} 与 NumPy 共享内存: {shares_memory(arr, big)}")

big[0] = 42.0                                           # 修改 NumPy 源数组
print("修改源数组后：jnp.array[0] =", float(a_copy[0]), "| device_put[0] =", float(a_put[0]))
assert not shares_memory(a_copy, big) and float(a_copy[0]) == 1.0
assert not shares_memory(a_asarray, big)
assert shares_memory(a_put, big) and float(a_put[0]) == 42.0   # ⚠️ “不可变”的 JAX 数组被改了
assert shares_memory(a_put_noalias, big)                        # 0.12.0.1 实测：may_alias=False 也别名

# JAX → NumPy
y = jax.jit(lambda v: v * 2)(jnp.arange(1024.0))
n_view, n_copy = np.asarray(y), np.array(y)
print("np.asarray 零拷贝:", n_view.ctypes.data == y.unsafe_buffer_pointer(), "| 可写:", n_view.flags.writeable)
print("np.array   零拷贝:", n_copy.ctypes.data == y.unsafe_buffer_pointer(), "| 可写:", n_copy.flags.writeable)
assert n_view.ctypes.data == y.unsafe_buffer_pointer() and not n_view.flags.writeable
assert n_copy.ctypes.data != y.unsafe_buffer_pointer() and n_copy.flags.writeable

# 同步点的代价：每步都转回 NumPy，对比只在最后转一次（计时只打印，不做 assert）
@jax.jit
def step(x):
    return jnp.sin(x) * 0.999 + 0.001

x0 = jnp.ones((512, 512))

def loop_convert_each_step(x, n=100):
    hist = []
    for _ in range(n):
        x = step(x)
        hist.append(np.array(x)[0, 0])        # 每步都同步并复制
    return x

def loop_convert_at_end(x, n=100):
    for _ in range(n):
        x = step(x)                           # 异步排队，不阻塞
    return x

t_each = bench(loop_convert_each_step, x0, n=3)
t_end = bench(loop_convert_at_end, x0, n=3)
print(f"每步转换: {t_each:.2f} ms | 只在末尾转换: {t_end:.2f} ms | 比值 {t_each / t_end:.2f}x")
print("解读：" + ("每步同步更慢，符合预期。" if t_each > t_end * 1.1 else
      "本机仅 CPU：设备内存就是主机内存，“拷回”几乎免费，异步排队的收益也有限，两者差距不明显，甚至可能反转。"
      "同步点的代价主要出现在 GPU/TPU 上，那里每次都要等待设备并走 PCIe 传输。"))
np.testing.assert_allclose(loop_convert_each_step(x0), loop_convert_at_end(x0))

# 验证：在 jit 里对 tracer 调 NumPy 函数
@jax.jit
def bad(x):
    return np.sin(x)      # 追踪时 x 是 tracer，NumPy 无法把它转成具体数组

try:
    bad(jnp.ones(3))
    err = None
except Exception as e:
    err = e
    print(type(e).__name__, ":", str(e).splitlines()[0][:120])
assert isinstance(err, jax.errors.TracerArrayConversionError)

host = jax.device_get({"w": jnp.ones(2), "b": jnp.zeros(1)})
assert all(isinstance(v, np.ndarray) for v in host.values())

# ❌ 旧写法（过时/错误）：to_dlpack / device_buffer / from_dlpack(胶囊) / random.normal(device=) 在 0.12 均不可用
x_src = jax.random.normal(jax.random.key(0), (1024, 1024))
attempts = {
    "jax.dlpack.to_dlpack(x)": lambda: jax.dlpack.to_dlpack(x_src),
    "x.device_buffer.unsafe_buffer_pointer()": lambda: x_src.device_buffer.unsafe_buffer_pointer(),
    "jax.dlpack.from_dlpack(x.__dlpack__())": lambda: jax.dlpack.from_dlpack(x_src.__dlpack__()),
    "jax.random.normal(key, shape, device=...)": lambda: jax.random.normal(jax.random.key(0), (2,), device=jax.devices()[0]),
}
for desc, fn in attempts.items():
    try:
        fn(); print(f"{desc:45s} → 仍可用")
    except Exception as e:
        print(f"{desc:45s} → {type(e).__name__}: {str(e).splitlines()[0][:90]}")
print("jax.dlpack 公开成员:", [a for a in dir(jax.dlpack) if not a.startswith("_")])
print("（本环境没有安装 torch / cupy，这两段示例在这里跑不了）")

# ✅ 推荐写法：DLPack 协议 —— JAX → NumPy 零拷贝
x = jax.random.normal(jax.random.key(0), (1024, 1024))
x.block_until_ready()                       # 导出前先同步（注意事项 2）
print("__dlpack_device__:", x.__dlpack_device__())
print("__dlpack__() 返回:", type(x.__dlpack__()).__name__)

n = np.from_dlpack(x)                       # NumPy 自己调用 x.__dlpack__()
print("JAX 指针  :", x.unsafe_buffer_pointer())
print("NumPy 指针:", n.ctypes.data, "| 可写:", n.flags.writeable)
assert not hasattr(jax.dlpack, "to_dlpack")
assert n.ctypes.data == x.unsafe_buffer_pointer() == x.addressable_data(0).unsafe_buffer_pointer()
assert not n.flags.writeable                # 导出的是只读视图，保护 JAX 的不可变性
np.testing.assert_array_equal(n, x)

# ✅ 推荐写法：NumPy → JAX 零拷贝（jax.dlpack.from_dlpack / jnp.from_dlpack）
src = aligned_zeros(1 << 20); src[:] = np.arange(src.size, dtype=np.float32)
j1 = jax.dlpack.from_dlpack(src)
j2 = jnp.from_dlpack(src)                   # 等价入口
print("from_dlpack 共享内存:", shares_memory(j1, src), "| jnp.from_dlpack 共享内存:", shares_memory(j2, src))
assert shares_memory(j1, src) and shares_memory(j2, src)

# 对齐要求：起始地址不是 64 字节对齐时，默认（copy=None）会静默复制；copy=False 直接报错
unaligned = aligned_zeros(1025)[1:]         # 偏移 4 字节
assert unaligned.ctypes.data % 64 != 0
j_u = jax.dlpack.from_dlpack(unaligned)
print("未对齐数组 from_dlpack 共享内存:", shares_memory(j_u, unaligned))
assert not shares_memory(j_u, unaligned)
try:
    jax.dlpack.from_dlpack(unaligned, copy=False)
    raised = False
except ValueError as e:
    raised = True
    print("copy=False →", str(e)[:110])
assert raised

# 验证：零拷贝导入后修改源数组 → JAX 数组随之改变；生命周期由消费方持有引用保证
src = aligned_zeros(1024); src[:] = 1.0
j = jax.dlpack.from_dlpack(src)
before = float(j[0])
src[0] = 99.0
print("修改前 j[0] =", before, "| 修改源数组后 j[0] =", float(j[0]))
assert before == 1.0 and float(j[0]) == 99.0

# 生命周期：删掉 JAX 数组后，NumPy 视图仍然有效（DLPack 的 deleter 机制让消费方持有生产方的引用）
x = jnp.full((1000,), 3.0)
view = np.from_dlpack(x)
del x
import gc; gc.collect()
print("原 JAX 数组删除后，NumPy 视图求和 =", view.sum())
assert view.sum() == 3000.0

# 计时：零拷贝 vs 复制（64 MB；只打印比较结果）
big = aligned_zeros(16 << 20); big[:] = 1.0
t_copy = bench(lambda a: jnp.array(a), big, n=10)
t_dl = bench(lambda a: jax.dlpack.from_dlpack(a), big, n=10)
print(f"jnp.array(复制): {t_copy:.3f} ms | from_dlpack(零拷贝): {t_dl:.4f} ms | 快 {t_copy / t_dl:.0f}x")

# ❌ 旧写法（过时/错误）：from jax.experimental import host_callback —— 0.12 中该模块已删除
import importlib.util
print("find_spec:", importlib.util.find_spec("jax.experimental.host_callback"))
try:
    from jax.experimental import host_callback
    host_callback.id_tap(print, jnp.ones(2))
except ImportError as e:
    print("ImportError:", short_path(e))
assert importlib.util.find_spec("jax.experimental.host_callback") is None

# ✅ 推荐写法：id_tap → jax.debug.callback（对应课程代码块 1）
records = []
def log_intermediate_data(arg):
    records.append(type(arg))
    arg = np.asarray(arg)                    # 回调收到 jax.Array，先转 NumPy
    print(f"[主机回调] 形状={arg.shape} 前5个={np.round(arg.ravel()[:5], 4)}")

@jax.jit
def process_data(x):
    y = jnp.sin(x) * 2.0
    jax.debug.callback(log_intermediate_data, y)   # 不返回值，不影响数据流
    z = jnp.mean(jnp.cos(y))
    return z

input_data = jax.random.uniform(jax.random.key(42), (8, 8))
result = process_data(input_data)
result.block_until_ready()
jax.effects_barrier()                        # 等所有回调执行完，避免打印漏到别的单元
print("最终结果:", result)
assert len(records) == 1 and issubclass(records[0], jax.Array)
ref = np.mean(np.cos(np.sin(np.asarray(input_data, np.float64)) * 2.0))
np.testing.assert_allclose(result, ref, rtol=1e-5)

# ✅ 推荐写法：id_print → jax.debug.print（格式字符串，懒格式化）
@jax.jit
def simple_print_example(x):
    y = x + 5.0
    jax.debug.print("中间值 y = {y}", y=y)
    return y * 2.0

output = simple_print_example(jnp.arange(3.0))
output.block_until_ready(); jax.effects_barrier()
np.testing.assert_array_equal(output, [10.0, 12.0, 14.0])

# ✅ 推荐写法：host_callback.call（纯计算）→ jax.pure_callback（对应课程代码块 2）
def external_cpu_calculation(data, parameter):
    data = np.asarray(data)                  # 主机端：普通 NumPy 计算
    return (np.tanh(data) + np.asarray(parameter)).astype(data.dtype)

@jax.jit
def jax_workflow(x, ext_param):
    intermediate = jnp.log1p(jnp.abs(x))
    external_result = jax.pure_callback(
        external_cpu_calculation,
        jax.ShapeDtypeStruct(intermediate.shape, intermediate.dtype),   # 输出形状与 dtype
        intermediate, ext_param)
    return external_result / (1.0 + jnp.mean(intermediate))

input_array = jnp.linspace(-2.0, 2.0, 5, dtype=jnp.float32)
final_result = jax_workflow(input_array, 0.5)
print("最终结果:", final_result)
xi = np.log1p(np.abs(np.linspace(-2, 2, 5)))
np.testing.assert_allclose(final_result, (np.tanh(xi) + 0.5) / (1 + xi.mean()), rtol=1e-5)

# ✅ 推荐写法：有副作用的调用 → io_callback；ordered=True 保证执行顺序与程序顺序一致
from jax.experimental import io_callback

order = []
def write_log(i):
    order.append(int(i))                     # 副作用：写外部状态

def body(i, carry):
    io_callback(write_log, None, i, ordered=True)   # 返回 None 时 result_shape_dtypes 传 None
    return carry + i

total = jax.jit(lambda: jax.lax.fori_loop(0, 5, body, 0))()
jax.effects_barrier()
print("写入顺序:", order, "| total =", total)
assert order == [0, 1, 2, 3, 4] and int(total) == 10

# 逐项实测上表
def run(fn):
    try:
        jax.block_until_ready(fn()); jax.effects_barrier()
        return "OK"
    except Exception as e:
        return type(e).__name__

calls = []
def host_log(a):  calls.append(np.shape(a))
def host_double(a):
    calls.append(np.shape(a)); return np.asarray(a) * 2

sds = lambda v: jax.ShapeDtypeStruct(v.shape, v.dtype)
variants = {
    "debug.callback":           lambda v: (jax.debug.callback(host_log, v), v * 2)[1],
    "io_callback":              lambda v: io_callback(host_double, sds(v), v),
    "io_callback(ordered)":     lambda v: io_callback(host_double, sds(v), v, ordered=True),
    "pure_callback(默认)":      lambda v: jax.pure_callback(host_double, sds(v), v),
    "pure_callback(expand_dims)": lambda v: jax.pure_callback(host_double, sds(v), v, vmap_method="expand_dims"),
}
x3, xb = jnp.arange(3.0), jnp.ones((4, 3))
table = {}
for name, f in variants.items():
    row = {}
    row["jit"] = run(lambda: jax.jit(f)(x3))
    calls.clear(); row["vmap"] = run(lambda: jax.vmap(f)(xb)); row["vmap调用次数"] = len(calls)
    row["grad"] = run(lambda: jax.grad(lambda v: f(v).sum())(x3))
    calls.clear(); jax.jit(lambda v: (f(v), v + 1)[1])(x3).block_until_ready(); jax.effects_barrier()
    row["未使用时执行次数"] = len(calls)
    table[name] = row

print(f"{'回调':28s}{'jit':6s}{'vmap':22s}{'vmap调用次数':10s}{'grad':12s}{'未使用时执行次数'}")
for name, r in table.items():
    print(f"{name:28s}{r['jit']:6s}{r['vmap']:22s}{r['vmap调用次数']!s:14s}{r['grad']:12s}{r['未使用时执行次数']}")

assert all(r["jit"] == "OK" for r in table.values())
assert table["debug.callback"]["vmap"] == "OK" and table["debug.callback"]["vmap调用次数"] == 4
assert table["debug.callback"]["grad"] == "OK"
assert table["io_callback"]["vmap"] == "OK" and table["io_callback"]["grad"] == "ValueError"
assert table["io_callback(ordered)"]["vmap"] == "ValueError"
assert table["pure_callback(默认)"]["vmap"] == "NotImplementedError"
assert table["pure_callback(默认)"]["grad"] == "ValueError"
assert table["pure_callback(expand_dims)"]["vmap"] == "OK" and table["pure_callback(expand_dims)"]["vmap调用次数"] == 1
assert table["debug.callback"]["未使用时执行次数"] == 1 and table["io_callback"]["未使用时执行次数"] == 1
assert table["pure_callback(默认)"]["未使用时执行次数"] == 0      # 纯回调被当作死代码删除

# ✅ 推荐写法：课程主示例（jit 部分），key 用 jax.random.key 生成
def external_pure_python_computation(x, y):
    """纯 Python/NumPy 函数：确定性、无副作用"""
    x = np.asarray(x)                     # 回调收到 jax.Array，转 NumPy
    return (np.sin(x) * y + 1.0).astype(x.dtype)

def jax_function_using_callback(a, b):
    output_shape_dtype = jax.ShapeDtypeStruct(a.shape, a.dtype)
    result = jax.pure_callback(
        external_pure_python_computation, output_shape_dtype, a, b,
        vmap_method="expand_dims")        # 显式声明 vmap 语义（见下文对照）
    return result + a

key = jax.random.key(0)
input_array = jax.random.normal(key, (3, 3))
input_scalar = 2.0
output = jax.jit(jax_function_using_callback)(input_array, input_scalar)
expected = np.sin(np.asarray(input_array, np.float64)) * input_scalar + 1.0 + np.asarray(input_array, np.float64)
print("输出:\n", output)
np.testing.assert_allclose(output, expected, rtol=1e-5)

# vmap：expand_dims 下回调只调用一次，收到 (10, 3, 3) 的整批数据
batched_input = jax.random.normal(key, (10, 3, 3))
batched_output = jax.vmap(jax_function_using_callback, in_axes=(0, None))(batched_input, input_scalar)
print("批量输出形状:", batched_output.shape)
assert batched_output.shape == (10, 3, 3)
np.testing.assert_allclose(batched_output, np.sin(np.asarray(batched_input)) * 2 + 1 + np.asarray(batched_input), rtol=1e-5)

# ❌ 旧写法（过时/错误）：不指定 vmap_method 直接 vmap；直接对 pure_callback 求 grad
def course_version(a, b):
    out = jax.pure_callback(external_pure_python_computation, jax.ShapeDtypeStruct(a.shape, a.dtype), a, b)
    return out + a

try:
    jax.vmap(course_version, in_axes=(0, None))(batched_input, input_scalar)
except NotImplementedError as e:
    print("vmap → NotImplementedError:", str(e)[:150])
try:
    jax.grad(lambda x: jnp.sum(course_version(x, input_scalar)))(input_array)
except ValueError as e:
    print("grad → ValueError:", str(e)[:150])

# ✅ 推荐写法：custom_jvp 给 pure_callback 补上导数（d/dx [sin(x)·y + 1] = cos(x)·y）
@jax.custom_jvp
def ext_fn(a, b):
    return jax.pure_callback(external_pure_python_computation, jax.ShapeDtypeStruct(a.shape, a.dtype),
                             a, b, vmap_method="expand_dims")

@ext_fn.defjvp
def ext_fn_jvp(primals, tangents):
    a, b = primals
    a_dot, b_dot = tangents
    y = ext_fn(a, b)
    # 导数部分用 jnp 写，它本身可微、可 jit、可 vmap
    return y, jnp.cos(a) * b * a_dot + jnp.sin(a) * b_dot

grad_fn = jax.jit(jax.grad(lambda x: jnp.sum(ext_fn(x, input_scalar) + x)))
gradients = grad_fn(input_array)
print("梯度:\n", gradients)
np.testing.assert_allclose(gradients, np.cos(np.asarray(input_array)) * 2.0 + 1.0, rtol=1e-5)
# 对标量参数 b 求导
gb = jax.grad(lambda b: jnp.sum(ext_fn(input_array, b)))(2.0)
np.testing.assert_allclose(gb, np.sin(np.asarray(input_array)).sum(), rtol=1e-5)
# vmap ∘ grad 也能组合
vg = jax.vmap(jax.grad(lambda x: jnp.sum(ext_fn(x, 2.0))))(batched_input)
np.testing.assert_allclose(vg, np.cos(np.asarray(batched_input)) * 2.0, rtol=1e-5)

# 课程多输出示例：result_shape_dtypes 用元组（pytree）描述
def multi_output_pure_function(x):
    x = np.asarray(x)
    return np.sum(x).astype(x.dtype), np.mean(x).astype(x.dtype)

def jax_multi_output_callback(arr):
    output_structure = (jax.ShapeDtypeStruct((), arr.dtype),   # 标量形状为 ()
                        jax.ShapeDtypeStruct((), arr.dtype))
    sum_val, mean_val = jax.pure_callback(multi_output_pure_function, output_structure, arr)
    return sum_val * 2, mean_val * 3

res1, res2 = jax.jit(jax_multi_output_callback)(jnp.arange(5.0))
print(f"多个输出：{res1=}, {res2=}")
assert float(res1) == 20.0 and float(res2) == 6.0

# 逐个试一遍 vmap_method：调用次数与回调收到的形状
seen = []
def h2(a, b):
    seen.append((np.shape(a), np.shape(b)))
    return (np.asarray(a) * np.asarray(b)).astype(np.float32)

def call_with(method):
    def f(a):
        return jax.pure_callback(h2, jax.ShapeDtypeStruct(a.shape, a.dtype), a, jnp.ones(3), vmap_method=method)
    return f

results = {}
for m in ["sequential", "sequential_unrolled", "expand_dims", "broadcast_all", "legacy_vectorized"]:
    seen.clear()
    out = jax.vmap(call_with(m))(jnp.ones((4, 3)))
    results[m] = list(seen)
    print(f"{m:20s} 调用 {len(seen)} 次，收到 (a, b) 形状 = {seen[0]}  输出 {out.shape}")

assert len(results["sequential"]) == 4 and results["sequential"][0] == ((3,), (3,))
assert results["expand_dims"] == [((4, 3), (1, 3))]       # 未批处理的 b 只多一个 1 维
assert results["broadcast_all"] == [((4, 3), (4, 3))]     # b 被广播到批大小
assert results["legacy_vectorized"] == [((4, 3), (3,))]   # b 保持原形状
try:
    jax.vmap(call_with("bogus"))(jnp.ones((4, 3)))
except ValueError as e:
    print("非法取值 →", str(e)[:120])

# 验证：导数本身也用 pure_callback 时，hessian 会失败；把导数也包成 custom_jvp 即可修复
np_cos_raw = lambda x: jax.pure_callback(lambda a: np.cos(np.asarray(a)).astype(a.dtype),
                                          jax.ShapeDtypeStruct(x.shape, x.dtype), x, vmap_method="expand_dims")
@jax.custom_jvp
def np_sin_v1(x):
    return jax.pure_callback(lambda a: np.sin(np.asarray(a)).astype(a.dtype),
                             jax.ShapeDtypeStruct(x.shape, x.dtype), x, vmap_method="expand_dims")
np_sin_v1.defjvp(lambda p, t: (np_sin_v1(p[0]), np_cos_raw(p[0]) * t[0]))

xs = jnp.linspace(0.0, 1.0, 4)
np.testing.assert_allclose(jax.grad(lambda v: np_sin_v1(v).sum())(xs), np.cos(np.asarray(xs)), rtol=1e-6)
try:
    jax.hessian(lambda v: np_sin_v1(v).sum())(xs)
    ok = True
except ValueError as e:
    ok = False
    print("hessian(v1) → ValueError:", str(e)[:80])
assert not ok

# 修复：sin 和 cos 互为导数，两者都定义 custom_jvp
@jax.custom_jvp
def np_sin(x):
    return jax.pure_callback(lambda a: np.sin(np.asarray(a)).astype(a.dtype), jax.ShapeDtypeStruct(x.shape, x.dtype), x, vmap_method="expand_dims")
@jax.custom_jvp
def np_cos(x):
    return np_cos_raw(x)
np_sin.defjvp(lambda p, t: (np_sin(p[0]), np_cos(p[0]) * t[0]))
np_cos.defjvp(lambda p, t: (np_cos(p[0]), -np_sin(p[0]) * t[0]))
H = jax.hessian(lambda v: np_sin(v).sum())(xs)
np.testing.assert_allclose(np.diag(H), -np.sin(np.asarray(xs)), rtol=1e-5)
print("修复后 hessian 对角线:", np.diag(H))

# 验证：形状不匹配会在运行时报错；输出未被使用时回调不会执行
try:
    jax.jit(lambda v: jax.pure_callback(lambda a: np.asarray(a)[:2], jax.ShapeDtypeStruct(v.shape, v.dtype), v))(jnp.arange(3.0)).block_until_ready()
    shape_err = None
except Exception as e:
    shape_err = type(e).__name__
print("返回形状不符 →", shape_err)
assert shape_err == "JaxRuntimeError"

counter = []
def impure(a):
    counter.append(1)                         # ❌ 在 pure_callback 里做副作用（反例）
    return np.asarray(a)
out = jax.jit(lambda v: (jax.pure_callback(impure, jax.ShapeDtypeStruct(v.shape, v.dtype), v), v + 1)[1])(jnp.arange(3.0))
out.block_until_ready()
print("输出未被使用时回调执行次数:", len(counter))
assert len(counter) == 0

# 计时：pure_callback（NumPy sin）vs 原生 jnp.sin（只打印）
xbig = jax.random.normal(jax.random.key(1), (1000, 1000))
f_cb = jax.jit(lambda v: jax.pure_callback(lambda a: np.sin(np.asarray(a)), jax.ShapeDtypeStruct(v.shape, v.dtype), v) * 2 + 1)
f_native = jax.jit(lambda v: jnp.sin(v) * 2 + 1)
np.testing.assert_allclose(f_cb(xbig), f_native(xbig), rtol=1e-5, atol=1e-6)
t_cb, t_nat = bench(f_cb, xbig), bench(f_native, xbig)
print(f"pure_callback: {t_cb:.3f} ms | 原生 jnp: {t_nat:.3f} ms | 比值 {t_cb / t_nat:.1f}x")

# 查看一个函数分解成哪些原语
def f(x, y):
    return jnp.sum(jnp.sin(x) * y + jnp.dot(x, y))

closed = jax.make_jaxpr(f)(jnp.ones(3), jnp.ones(3))
print(closed)
names = [eqn.primitive.name for eqn in closed.jaxpr.eqns]
print("原语序列:", names)
assert {"sin", "mul", "dot_general", "add", "reduce_sum"} <= set(names)

# 控制流原语的名字
def g(x):
    y = jax.lax.cond(x.sum() > 0, jnp.sin, jnp.cos, x)
    return jax.lax.scan(lambda c, v: (c + v, c), 0.0, y)[0]
cf_names = {e.primitive.name for e in jax.make_jaxpr(g)(jnp.ones(3)).jaxpr.eqns}
print("控制流原语:", cf_names & {"cond", "scan"})
assert {"cond", "scan"} <= cf_names

import jax.extend.core.primitives as P          # 内置原语对象的半公开入口
print("P.scan_p.name =", P.scan_p.name, "| P.cond_p.name =", P.cond_p.name, "| jax.lax.add_p.name =", jax.lax.add_p.name)
assert (P.scan_p.name, P.cond_p.name, jax.lax.add_p.name) == ("scan", "cond", "add")

# 原语身上挂着各种规则表：JVP、转置、批处理、lowering
from jax.interpreters import ad, batching, mlir
print("sin 有 JVP 规则     :", P.sin_p in ad.primitive_jvps)
print("mul 有转置规则      :", P.mul_p in ad.primitive_transposes)
# 注意：0.12 的 batching.primitive_batchers 是“只写代理”（只支持 [p] = rule），不能用 in 查询；
# 它把规则包装后存进 batching.fancy_primitive_batchers
try:
    P.sin_p in batching.primitive_batchers
except TypeError as e:
    print("primitive_batchers 不支持 in 查询:", e)
print("sin 有 batching 规则:", P.sin_p in batching.fancy_primitive_batchers)
assert P.sin_p in ad.primitive_jvps and P.mul_p in ad.primitive_transposes and P.sin_p in batching.fancy_primitive_batchers

# 编译流水线：jaxpr → StableHLO（MLIR 文本）
hlo_text = jax.jit(f).lower(jnp.ones(3), jnp.ones(3)).as_text()
print(hlo_text)
assert "stablehlo.sine" in hlo_text and "stablehlo.dot_general" in hlo_text

# ❌ 旧写法（过时/错误）：from jax.core import Primitive —— 0.12 中 jax.core 不再导出 Primitive
try:
    from jax.core import Primitive as _P
except ImportError as e:
    print("ImportError:", short_path(e))
import jax.core, jax.extend.core
print("jax.core 有 Primitive:", hasattr(jax.core, "Primitive"), "| 有 ShapedArray:", hasattr(jax.core, "ShapedArray"))
print("jax.extend.core 有 Primitive:", hasattr(jax.extend.core, "Primitive"), "| 有 ShapedArray:", hasattr(jax.extend.core, "ShapedArray"))

# ✅ 推荐写法：步骤 1～3 —— 原语对象、用户函数、impl 与抽象求值
from jax.extend.core import Primitive          # 半内部 API：0.12 下验证过的路径
from jax.core import ShapedArray               # 抽象值类型
from jax.interpreters import mlir, ad, batching

my_custom_op_p = Primitive("my_custom_op")      # 步骤 1：唯一名字

def my_custom_op(x, *, some_parameter):         # 步骤 2：面向用户的函数
    return my_custom_op_p.bind(x, some_parameter=some_parameter)

@my_custom_op_p.def_impl                         # eager 实现：收到具体数组
def _my_custom_op_impl(x, *, some_parameter):
    return x * x + some_parameter

@my_custom_op_p.def_abstract_eval                # 步骤 3：只看形状/dtype，不算数值
def _my_custom_op_abstract_eval(abstract_x, *, some_parameter):
    if some_parameter < 0:
        raise ValueError("some_parameter must be non-negative")
    return ShapedArray(abstract_x.shape, abstract_x.dtype)

x = jnp.array([1.0, 2.0, 3.0])
print("eager:", my_custom_op(x, some_parameter=1.0))
print(jax.make_jaxpr(partial(my_custom_op, some_parameter=1.0))(x))
np.testing.assert_array_equal(my_custom_op(x, some_parameter=1.0), [2.0, 5.0, 10.0])

# 还没有 lowering → jit 失败；没有 JVP → grad 失败；没有 batching → vmap 失败
for name, fn in [("jit", lambda: jax.jit(partial(my_custom_op, some_parameter=1.0))(x)),
                 ("grad", lambda: jax.grad(lambda v: my_custom_op(v, some_parameter=1.0).sum())(x)),
                 ("vmap", lambda: jax.vmap(partial(my_custom_op, some_parameter=1.0))(jnp.ones((2, 3))))]:
    try:
        fn(); print(name, "→ 意外成功")
    except NotImplementedError as e:
        print(f"{name:5s}→ NotImplementedError: {str(e).splitlines()[0][:80]}")

# ✅ 推荐写法：步骤 4 —— lowering（详见 5.8；这里用 mlir.lower_fun 把 jnp 实现降低成 StableHLO）
mlir.register_lowering(
    my_custom_op_p,
    mlir.lower_fun(lambda x, *, some_parameter: x * x + some_parameter, multiple_results=False))

y_jit = jax.jit(partial(my_custom_op, some_parameter=1.0))(x)
print("jit:", y_jit)
np.testing.assert_array_equal(y_jit, [2.0, 5.0, 10.0])

# ❌ 旧写法（过时/错误）：JVP 规则里用 ad.Zero.from_value(primal_out) 构造零切线 —— 0.12 中该方法不存在
print("ad.Zero 有 from_value:", hasattr(ad.Zero, "from_value"))
try:
    ad.Zero.from_value(jnp.ones(3))
except AttributeError as e:
    print("AttributeError:", e)

# ✅ 推荐写法：步骤 5 —— JVP 规则（ẏ = 2x·ẋ），零切线用 ad.Zero(切线 aval) 构造
def _my_custom_op_jvp(primals, tangents, *, some_parameter):
    x, = primals
    x_dot, = tangents
    primal_out = my_custom_op_p.bind(x, some_parameter=some_parameter)
    if type(x_dot) is ad.Zero:                                  # 输入被当作常数
        tangent_out = ad.Zero(jax.typeof(primal_out).to_tangent_aval())
    else:
        tangent_out = 2 * x * x_dot                             # 用现有原语表达，对 x_dot 是线性的
    return primal_out, tangent_out

ad.primitive_jvps[my_custom_op_p] = _my_custom_op_jvp

g = jax.grad(lambda v: my_custom_op(v, some_parameter=1.0).sum())(x)
gj = jax.jit(jax.grad(lambda v: my_custom_op(v, some_parameter=1.0).sum()))(x)
print("grad:", g, "| jit(grad):", gj)
np.testing.assert_allclose(g, 2 * np.asarray(x)); np.testing.assert_allclose(gj, 2 * np.asarray(x))

# 逐元素原语的 batching：batching.defvectorized 一行搞定（batch 维原样透传）
batching.defvectorized(my_custom_op_p)
xb = jnp.arange(6.0).reshape(2, 3)
yb = jax.vmap(partial(my_custom_op, some_parameter=1.0))(xb)
np.testing.assert_array_equal(yb, np.asarray(xb) ** 2 + 1)
vg = jax.vmap(jax.grad(lambda v: my_custom_op(v, some_parameter=0.0).sum()))(xb)
np.testing.assert_array_equal(vg, 2 * np.asarray(xb))
print("vmap / vmap(grad) 通过")

# ✅ 课程示例 1：custom_add_one 的抽象求值（含输入校验）
custom_add_one_p = Primitive("custom_add_one")

@custom_add_one_p.def_abstract_eval
def custom_add_one_abstract_eval(*avals, **params):
    if len(avals) != 1:
        raise ValueError("custom_add_one 需要恰好一个输入操作数。")
    input_aval = avals[0]
    if not isinstance(input_aval, ShapedArray):
        raise TypeError(f"期望 ShapedArray，得到 {type(input_aval)}")
    if not jnp.issubdtype(input_aval.dtype, jnp.floating):
        raise TypeError(f"custom_add_one 只支持浮点输入，得到 {input_aval.dtype}")
    return ShapedArray(input_aval.shape, input_aval.dtype)

impl_calls = []
@custom_add_one_p.def_impl
def _add_one_impl(x):
    impl_calls.append(1)
    return x + 1.0

# jax.eval_shape 只做抽象求值：得到形状，但 impl 一次都不执行
out_sds = jax.eval_shape(custom_add_one_p.bind, jax.ShapeDtypeStruct((4, 5), jnp.float32))
print("eval_shape:", out_sds, "| impl 调用次数:", len(impl_calls))
assert out_sds.shape == (4, 5) and out_sds.dtype == jnp.float32 and impl_calls == []

# 类型错误在追踪期就被发现（还没编译、没执行）
try:
    jax.eval_shape(custom_add_one_p.bind, jax.ShapeDtypeStruct((3,), jnp.int32))
    caught = False
except TypeError as e:
    caught = True
    print("追踪期 TypeError:", e)
assert caught

# ✅ 课程示例 2：custom_reduce_sum —— 输出形状由参数 axis 决定
custom_reduce_sum_p = Primitive("custom_reduce_sum")

def custom_reduce_sum(x, axis):
    axis = tuple(sorted(a % x.ndim for a in (axis if isinstance(axis, tuple) else (axis,))))
    return custom_reduce_sum_p.bind(x, axis=axis)         # axis 必须可哈希（tuple）

@custom_reduce_sum_p.def_abstract_eval
def custom_reduce_sum_abstract_eval(aval, *, axis):
    if not isinstance(aval, ShapedArray):
        raise TypeError("输入必须是 ShapedArray")
    output_shape = tuple(d for i, d in enumerate(aval.shape) if i not in axis)
    return ShapedArray(output_shape, aval.dtype)

custom_reduce_sum_p.def_impl(lambda x, *, axis: jnp.sum(x, axis=axis))
mlir.register_lowering(custom_reduce_sum_p,
                       mlir.lower_fun(lambda x, *, axis: jnp.sum(x, axis=axis), multiple_results=False))

my_array = jnp.arange(24.0).reshape(2, 3, 4)
print(jax.eval_shape(lambda a: custom_reduce_sum(a, (0,)), my_array))
assert jax.eval_shape(lambda a: custom_reduce_sum(a, (0, 2)), my_array).shape == (3,)
r = jax.jit(lambda a: custom_reduce_sum(a, (0,)))(my_array)
np.testing.assert_array_equal(r, np.asarray(my_array).sum(0))
print(jax.make_jaxpr(lambda a: custom_reduce_sum(a, -1))(my_array))

# 多输出原语：minmax(x) -> (min, max)
minmax_bad = Primitive("minmax_bad")                    # 忘了设 multiple_results
minmax_bad.def_abstract_eval(lambda a: (ShapedArray((), a.dtype), ShapedArray((), a.dtype)))
try:
    jax.eval_shape(minmax_bad.bind, jnp.arange(5.0))
    err = None
except ValueError as e:
    err = e; print("ValueError:", str(e)[:110])
assert err is not None

minmax_p = Primitive("minmax")
minmax_p.multiple_results = True
minmax_p.def_impl(lambda x: (jnp.min(x), jnp.max(x)))
minmax_p.def_abstract_eval(lambda a: (ShapedArray((), a.dtype), ShapedArray((), a.dtype)))
mlir.register_lowering(minmax_p, mlir.lower_fun(lambda x: (jnp.min(x), jnp.max(x)), multiple_results=True))
lo, hi = jax.jit(minmax_p.bind)(jnp.array([3.0, -1.0, 7.0]))
print("min, max =", lo, hi)
assert (float(lo), float(hi)) == (-1.0, 7.0)