第 2 章　用 JIT 编译加速函数

# 公共导入与工具函数（后面所有单元都依赖这里）
import time
import timeit
import logging
import re
import contextlib
import warnings
from functools import partial

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax

print("JAX 版本:", jax.__version__)
print("设备列表:", jax.devices())
print("数组所在设备:", jnp.ones(1).device)   # 0.12 推荐用 x.device 属性


def bench(f, *args, warmup=2, runs=20, **kwargs):
    """计时工具：先预热（触发编译），再多次运行取平均。

    用 jax.block_until_ready 等待结果真正算完——它能处理任意 PyTree，
    也能处理非 JAX 返回值（原样返回）。返回 (平均秒数, 标准差)。
    """
    for _ in range(warmup):
        jax.block_until_ready(f(*args, **kwargs))
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args, **kwargs))
        ts.append(time.perf_counter() - t0)
    return float(np.mean(ts)), float(np.std(ts))


def assert_faster(fast, slow, factor=1.0, tries=3):
    """计时断言：fast() 的耗时 × factor < slow() 的耗时（两个参数都是返回耗时的无参函数）。
    计时受机器负载影响（例如 CI 上几个 notebook 并行），不满足时重新测量，最多 tries 次。"""
    for _ in range(tries):
        tf, ts = fast(), slow()
        if tf * factor < ts:
            return tf, ts
    raise AssertionError(f"重测 {tries} 次仍不满足：{tf:.3g} × {factor} ≥ {ts:.3g}")

def expect_error(fn, exc_type=Exception):
    """断言 fn() 一定抛出 exc_type：打印异常类型和第一行信息，并返回异常对象。
    若没有报错则抛出 AssertionError——用来在 notebook 里“自证”某段代码确实会失败。"""
    try:
        out = fn()
    except exc_type as e:
        first = str(e).strip().splitlines()[0]
        print(f"❌ {type(e).__name__}: {first[:160]}")
        return e
    raise AssertionError(f"预期抛出 {exc_type.__name__}，实际没有报错，结果: {out}")


class _CompileLog(logging.Handler):
    """收集 jax_log_compiles 产生的 “Compiling jit(<函数名>)” 日志。"""

    def __init__(self, echo):
        super().__init__()
        self.names, self.echo = [], echo

    def emit(self, record):
        m = re.match(r"Compiling jit\((.*?)\)", record.getMessage())
        if m:
            self.names.append(m.group(1))
            if self.echo:
                print(f"    [编译] jit({m.group(1)})", flush=True)

    def count(self, name=None):
        return sum(1 for n in self.names if name is None or n == name)


@contextlib.contextmanager
def count_compiles(echo=True):
    """上下文管理器：统计块内真正发生的 XLA 编译次数（可按函数名过滤）。

    原理：临时打开 jax_log_compiles，并在 "jax" logger 上挂一个收集器。
    用法：with count_compiles() as log: ...;  assert log.count("函数名") == k
    """
    h = _CompileLog(echo)
    logger = logging.getLogger("jax")
    logger.addHandler(h)
    old = jax.config.jax_log_compiles
    jax.config.update("jax_log_compiles", True)
    try:
        yield h
    finally:
        jax.config.update("jax_log_compiles", old)
        logger.removeHandler(h)

def many_small_ops(x):
    # 50 个微小的逐元素操作：每个都要单独调度一次
    for _ in range(50):
        x = x * 1.0001 + 0.001
    return x

x_small = jnp.ones(100)
many_small_ops_jit = jax.jit(many_small_ops)

t_eager, _ = bench(many_small_ops, x_small)
t_jit, _ = bench(many_small_ops_jit, x_small)
print(f"eager: {t_eager*1e6:8.1f} µs")
print(f"jit  : {t_jit*1e6:8.1f} µs   加速约 {t_eager/t_jit:.0f} 倍")
# jit 版本必然更快（这里主要省掉的是 50 次调度开销）；计时受负载影响，assert_faster 不满足时会重测
assert_faster(lambda: bench(many_small_ops_jit, x_small)[0], lambda: bench(many_small_ops, x_small)[0])

# 查看 XLA 编译后的 HLO，寻找 fusion（融合内核）
fused = jax.jit(lambda a, x, b: jnp.tanh(a * x + b))
hlo = fused.lower(2.0, jnp.ones(1000), 1.0).compile().as_text()
fusion_lines = [line.strip() for line in hlo.splitlines() if "fusion(" in line]
for line in fusion_lines:
    print(line[:160])
assert len(fusion_lines) == 1   # 乘、加、tanh 只剩一条 fusion 指令
# 只有一条 fusion 指令：乘、加、tanh 被合并为一个循环内核（kind=kLoop）

# 课程示例：一个“稍复杂”的计算
def complex_computation(x, weight, bias):
    y = jnp.dot(x, weight) + bias
    z = jnp.tanh(y)
    return jnp.mean(z)

# 生成数据：用 jax.random.key 并为每个数组 split 出独立子 key（新旧写法的对照见本节后面）
key = jax.random.key(0)
kx, kw, kb = jax.random.split(key, 3)
x_data = jax.random.normal(kx, (1000, 500))
weight_data = jax.random.normal(kw, (500, 200))
bias_data = jax.random.normal(kb, (200,))

# 用法 1：装饰器
@jax.jit
def compiled_computation_decorator(x, weight, bias):
    y = jnp.dot(x, weight) + bias
    z = jnp.tanh(y)
    return jnp.mean(z)

# 用法 2：高阶函数
compiled_computation_functional = jax.jit(complex_computation)

# 计时流程（用更精确的 time.perf_counter，并统一调用 block_until_ready）

# 原始函数：先预热一次，再计时
_ = complex_computation(x_data, weight_data, bias_data).block_until_ready()
t0 = time.perf_counter()
result_original = complex_computation(x_data, weight_data, bias_data).block_until_ready()
print(f"原始函数时间:                      {time.perf_counter() - t0:.6f} 秒")

# 装饰器版本：第一次调用包含 追踪 + 编译
t0 = time.perf_counter()
result_compiled_decorator = compiled_computation_decorator(x_data, weight_data, bias_data).block_until_ready()
print(f"编译函数（装饰器）首次调用（含编译）: {time.perf_counter() - t0:.6f} 秒")

# 第二次调用：命中缓存
t0 = time.perf_counter()
result_compiled_decorator = compiled_computation_decorator(x_data, weight_data, bias_data).block_until_ready()
print(f"编译函数（装饰器）第二次调用（缓存）: {time.perf_counter() - t0:.6f} 秒")

print("结果匹配:", jnp.allclose(result_original, result_compiled_decorator))
assert jnp.allclose(result_original, result_compiled_decorator)

# 函数式版本：先调用一次触发编译，再计时
_ = compiled_computation_functional(x_data, weight_data, bias_data).block_until_ready()
t0 = time.perf_counter()
_ = compiled_computation_functional(x_data, weight_data, bias_data).block_until_ready()
print(f"编译函数（函数式）后续调用:          {time.perf_counter() - t0:.6f} 秒")

# ❌ 旧写法（过时）：PRNGKey 返回无类型的 uint32 原始 key；且同一个 key 被复用于多个数组
try:
    key_old = jax.random.PRNGKey(0)
    a_old = jax.random.normal(key_old, (6,))
    b_old = jax.random.normal(key_old, (3,))
    print("key =", key_old, " dtype =", key_old.dtype)
    print("a[:3] =", a_old[:3])
    print("b     =", b_old)
    print("两个数组前缀相同:", bool(jnp.allclose(a_old[:3], b_old)))
except Exception as e:
    print("报错:", type(e).__name__, e)

# ✅ 推荐写法：jax.random.key 返回带类型的 key；每个用途 split 出独立的子 key
key = jax.random.key(0)
k_a, k_b = jax.random.split(key)
a_new = jax.random.normal(k_a, (6,))
b_new = jax.random.normal(k_b, (3,))
print("key dtype =", key.dtype)
print("a[:3] =", a_new[:3])
print("b     =", b_new)

assert jnp.issubdtype(key.dtype, jax.dtypes.prng_key)     # 新式 key 是专门的 key 类型
assert not jnp.allclose(a_new[:3], b_new)                 # 独立子 key → 序列不再相同
# 反证：复用同一个 key，前缀必然相同（这正是要避免的）
assert jnp.allclose(jax.random.normal(key, (6,))[:3], jax.random.normal(key, (3,)))
print("✅ 断言通过")

def double(x):
    return x * 2

# 同一个函数对象 → 共享编译缓存
a_jit = jax.jit(double)
b_jit = jax.jit(double)
a_jit(1.0).block_until_ready()
t0 = time.perf_counter(); b_jit(1.0).block_until_ready()
print(f"另一个 jax.jit(double) 首次调用: {(time.perf_counter()-t0)*1e3:.3f} ms（命中共享缓存）")

# 循环里每次新建 lambda → 每次都重新编译
for i in range(3):
    t0 = time.perf_counter()
    jax.jit(lambda x: x * 2)(1.0).block_until_ready()
    print(f"第 {i} 次新 lambda: {(time.perf_counter()-t0)*1e3:.3f} ms（每次都在编译）")

# partial 写法传 jit 参数
@partial(jax.jit, static_argnames=["n"])
def power(x, n):
    return x ** n

print("power(2.0, n=3) =", power(2.0, n=3))
assert power(2.0, n=3) == 8.0

def my_simple_func(x, y):
    a = jnp.sin(x)
    b = jnp.cos(y)
    return a + b

x_example = jnp.ones(3)
y_example = jnp.zeros(3)

print(jax.make_jaxpr(my_simple_func)(x_example, y_example))

# 输入的数值不会出现在 Jaxpr 里；代码里写死的常量会以字面量出现
jp = str(jax.make_jaxpr(lambda x: x * 2.0)(jnp.array([7.5, 8.5])))
assert "2.0" in jp and "7.5" not in jp and "8.5" not in jp

@jax.jit
def slow_function(x):
    # 循环 5 次（Python 循环在追踪时会被展开）
    for _ in range(5):
        if x.shape[0] == x.shape[1]:
            # 方阵：矩阵乘法
            x = jnp.dot(x, x.T) + 0.5 * x
        else:
            # 非方阵：逐元素运算
            x = x * x + 0.5 * x
    return x

key = jax.random.key(0)
data = jax.random.normal(key, (100, 100))

t0 = time.perf_counter(); result1 = slow_function(data).block_until_ready()
print(f"首次调用（编译）耗时：    {time.perf_counter() - t0:.4f} 秒")

t0 = time.perf_counter(); result2 = slow_function(data).block_until_ready()
print(f"第二次调用（缓存）耗时：  {time.perf_counter() - t0:.4f} 秒")

data_different_shape = jax.random.normal(key, (150, 150))
t0 = time.perf_counter(); result3 = slow_function(data_different_shape).block_until_ready()
print(f"使用新形状调用耗时：      {time.perf_counter() - t0:.4f} 秒")

t0 = time.perf_counter(); result4 = slow_function(data).block_until_ready()
print(f"再次使用原始形状调用耗时：{time.perf_counter() - t0:.4f} 秒")

jaxpr = jax.make_jaxpr(slow_function.__wrapped__)(data)  # __wrapped__ 取回未 jit 的原函数
print("dot_general 出现次数:", str(jaxpr).count("dot_general"))
assert str(jaxpr).count("dot_general") == 5   # range(5) 被展开成 5 份

# 原始编译日志：先把输入建好——jnp.ones 本身在 eager 模式下也会编译小内核，会混进日志
a3, b3, a4 = jnp.ones(3), jnp.ones(3) * 7, jnp.ones(4)
g = jax.jit(lambda x: x * 2)
jax.config.update("jax_log_compiles", True)
print("调用 1：f32[3]", flush=True); g(a3)               # 编译
print("调用 2：f32[3] 不同数值", flush=True); g(b3)      # 命中缓存，没有日志
print("调用 3：f32[4]", flush=True); g(a4)               # 新形状 → 编译
jax.config.update("jax_log_compiles", False)

# E1 形状：只改数值不重新编译；改形状重新编译
def e1(x):
    return x * 2

x3a, x3b, x4 = jnp.ones(3), jnp.arange(3.0), jnp.ones(4)   # 先建好输入
e1_jit = jax.jit(e1)
with count_compiles() as log:
    e1_jit(x3a)
    e1_jit(x3b)                 # 同形状、不同数值
    n_same_shape = log.count("e1")
    e1_jit(x4)                  # 新形状
assert n_same_shape == 1
assert log.count("e1") == 2
print("E1 ✅ 同形状只编译 1 次；换形状后共编译", log.count("e1"), "次")

# E2 dtype：int32 会重新编译；默认配置下 float64 被截断为 float32，不会重新编译
def e2(x):
    return x + 1

xf32 = jnp.ones(3, dtype=jnp.float32)
xi32 = jnp.ones(3, dtype=jnp.int32)
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    x64_req = jnp.ones(3, dtype=jnp.float64)     # x64 默认关闭：被截断
assert x64_req.dtype == jnp.float32 and len(w) >= 1
print("请求 float64，实际得到:", x64_req.dtype, "| 警告:", str(w[0].message)[:60], "...")

e2_jit = jax.jit(e2)
with count_compiles() as log:
    e2_jit(xf32)
    e2_jit(x64_req)
    n_after_f64 = log.count("e2")
    e2_jit(xi32)
assert n_after_f64 == 1        # float32 与“float64”共用一份编译结果
assert log.count("e2") == 2    # int32 才触发第二次编译

# 对照：用上下文管理器临时开启 x64，float64 才是真正不同的 dtype（不影响后续单元）
def e2_x64(x):
    return x + 1

with jax.enable_x64(True):
    y32 = jnp.ones(3, dtype=jnp.float32)
    y64 = jnp.ones(3, dtype=jnp.float64)
    assert y64.dtype == jnp.float64
    e2_x64_jit = jax.jit(e2_x64)
    with count_compiles() as log64:
        e2_x64_jit(y32)
        e2_x64_jit(y64)
assert log64.count("e2_x64") == 2
print("E2 ✅ 默认配置：float32/“float64” 共 1 次编译，加 int32 共 2 次；开启 x64 后 float32/float64 各编译 1 次")

# E3 弱类型：Python 标量是“弱类型”，与同 dtype 的强类型数组签名不同
def e3(x):
    return x + 1

e3_jit = jax.jit(e3)
x_strong = jnp.float32(1.0)
with count_compiles() as log:
    e3_jit(1.0)
    e3_jit(2.0)          # Python float：弱类型 f32[]，只改数值 → 命中
    e3_jit(x_strong)     # 强类型 f32[] → 新签名
    e3_jit(1)            # Python int：弱类型 i32[] → 新签名
assert log.count("e3") == 3

# 由 Python 标量“派生”出来的数组往往也是弱类型
for name, arr in [("jnp.array(5.0)", jnp.array(5.0)), ("jnp.full(3, 7.0)", jnp.full(3, 7.0)),
                  ("jnp.ones(3)", jnp.ones(3)), ("jnp.float32(1.0)", jnp.float32(1.0))]:
    print(f"{name:18s} weak_type={arr.weak_type}")
assert jnp.array(5.0).weak_type and jnp.full(3, 7.0).weak_type
assert not jnp.ones(3).weak_type and not jnp.float32(1.0).weak_type
print("E3 ✅ 弱/强类型、float/int 标量共编译", log.count("e3"), "次")

# E4 静态参数值：每个新值编译一次，见过的值命中缓存
def e4(x, n):
    for _ in range(n):
        x = jnp.sin(x)
    return x

e4_jit = jax.jit(e4, static_argnames="n")
xe4 = jnp.ones(3)
with count_compiles() as log:
    for n in [2, 3, 2, 3]:
        e4_jit(xe4, n=n)
assert log.count("e4") == 2
print("E4 ✅ n 取 2,3,2,3 共编译", log.count("e4"), "次")

# E5 PyTree 结构：键、容器类型、长度都属于结构
def e5(d):
    return d["a"] * 2

def e5_seq(t):
    return t[0] * 2

d1 = {"a": jnp.ones(3)}
d2 = {"a": jnp.zeros(3)}                     # 同结构，不同数值
d3 = {"a": jnp.ones(3), "b": jnp.ones(3)}    # 多了一个键
xs5 = jnp.ones(3)
e5_jit, e5_seq_jit = jax.jit(e5), jax.jit(e5_seq)
with count_compiles() as log:
    e5_jit(d1)
    e5_jit(d2)
    n_same_tree = log.count("e5")
    e5_jit(d3)
    e5_seq_jit((xs5,))       # tuple
    e5_seq_jit([xs5])        # list：容器类型不同 → 新结构
assert n_same_tree == 1 and log.count("e5") == 2
assert log.count("e5_seq") == 2
print("E5 ✅ 字典加键 → 重新编译；tuple 与 list → 各编译一次")

# E6 位置 vs 关键字传参，以及省略默认参数
def e6(x, t):
    return x * t

e6_jit = jax.jit(e6, static_argnums=(1,))
x6 = jnp.ones(3)
with count_compiles() as log:
    e6_jit(x6, 2.0)          # 位置传入
    e6_jit(x6, t=2.0)        # 关键字传入：同样是静态，但缓存键不同
assert log.count("e6") == 2


def e6_default(x, tol=1e-6):
    return x + tol

e6d_jit = jax.jit(e6_default)
print(jax.make_jaxpr(e6_default)(x6))     # 省略时：默认值作为常量写进 Jaxpr（1e-6 的 float32 近似值显示为 9.99…e-07）
with count_compiles() as log:
    e6d_jit(x6)              # 省略：默认值被烘焙为常量
    e6d_jit(x6, 1e-3)        # 位置传入：tol 变成追踪参数 → 新条目
    e6d_jit(x6, 1e-4)        # 位置传入、换值 → 命中
    e6d_jit(x6, tol=1e-4)    # 关键字传入 → 新条目
assert log.count("e6_default") == 3
print("E6 ✅ 位置/关键字各占一个条目；省略/位置/关键字三种传法共编译 3 次")

# E7 每次新建 jit 对象：同一个函数对象共享缓存；每次新建 lambda 则每次都编译
def e7(x):
    return x * 2

x7 = jnp.ones(3)
with count_compiles() as log:
    for _ in range(3):
        jax.jit(e7)(x7)                  # 同一个函数对象 e7 → 共享缓存
    n_same_fn = log.count("e7")
    for _ in range(3):
        jax.jit(lambda x: x * 2)(x7)     # 每次都是新的 lambda 对象
assert n_same_fn == 1
assert log.count("<lambda>") == 3
print("E7 ✅ jax.jit(e7) ×3 只编译 1 次；新 lambda ×3 编译 3 次 → 在循环外 jit 一次")

# E8 全局变量：修改它不会触发重新编译 → 继续使用旧值
scale_e8 = 1.0

def e8(x):
    return x * scale_e8

e8_jit = jax.jit(e8)
x8 = jnp.ones(2)
with count_compiles() as log:
    r1 = e8_jit(x8)
    scale_e8 = 5.0               # 修改全局变量
    r2 = e8_jit(x8)
assert log.count("e8") == 1                 # 没有重新编译……
assert jnp.allclose(r2, 1.0)                # ……所以结果仍按旧值 1.0 计算
print("E8 ✅ 改全局变量后结果仍是", r2, "（更隐蔽的“新旧值并存”见 2.6 问题 5）")

def affine_sin(x):
    return jnp.sin(x) * 2

jf = jax.jit(affine_sin)
xs = jnp.ones(3)

print("== trace → Jaxpr ==")
print(jf.trace(xs).jaxpr)

lowered = jf.lower(xs)
print("== lower → StableHLO ==")
print(lowered.as_text())

t0 = time.perf_counter()
compiled = lowered.compile()
print(f"== compile 耗时 {(time.perf_counter()-t0)*1e3:.2f} ms ==")
print("执行结果:", compiled(xs))
assert jnp.allclose(compiled(xs), jnp.sin(xs) * 2)
ca = compiled.cost_analysis()
print("cost_analysis:", {k: ca[k] for k in ("flops", "transcendentals", "bytes accessed")})

# Compiled 对象不能处理不同形状
expect_error(lambda: compiled(jnp.ones(4)), TypeError)

def conditional_func_py(x):
    if x > 0:
        return x * 2
    else:
        return x / 2

print(conditional_func_py(5.0))    # 10.0
print(conditional_func_py(-4.0))   # -2.0

# 尝试 jit：x 被替换为追踪器，`if x > 0` 无法得到 True/False
jitted_conditional_py = jax.jit(conditional_func_py)
expect_error(lambda: jitted_conditional_py(jnp.array(5.0)), jax.errors.ConcretizationTypeError)

def conditional_func_lax(x):
    def true_branch(val):
        return val * 2

    def false_branch(val):
        return val / 2

    return lax.cond(x > 0, true_branch, false_branch, x)

jitted_conditional_lax = jax.jit(conditional_func_lax)
print(jitted_conditional_lax(jnp.array(5.0)))    # 10.0
print(jitted_conditional_lax(jnp.array(-4.0)))   # -2.0
assert jitted_conditional_lax(jnp.array(5.0)) == 10.0 and jitted_conditional_lax(jnp.array(-4.0)) == -2.0

def sum_first_n_py(arr, n):
    total = 0.0
    for i in range(n):
        total = total + arr[i]
    return total

jitted_sum_first_3 = jax.jit(sum_first_n_py, static_argnums=1)
my_array = jnp.arange(10.0)
print(jitted_sum_first_3(my_array, 3))   # 3.0 = 0 + 1 + 2
assert jitted_sum_first_3(my_array, 3) == 3.0

# 看看展开后的样子：n=3 时出现 3 次取元素 + 3 次加法
print(jax.make_jaxpr(sum_first_n_py, static_argnums=1)(my_array, 3))

def sum_first_n_lax(arr, n):
    def body_fun(i, current_total):
        return current_total + arr[i]

    initial_val = 0.0
    final_total = lax.fori_loop(0, n, body_fun, initial_val)
    return final_total

jitted_sum_lax = jax.jit(sum_first_n_lax)
my_array = jnp.arange(10.0)
print(jitted_sum_lax(my_array, jnp.array(3)))   # 3.0  = 0+1+2
print(jitted_sum_lax(my_array, jnp.array(5)))   # 10.0 = 0+1+2+3+4
assert jitted_sum_lax(my_array, jnp.array(3)) == 3.0 and jitted_sum_lax(my_array, jnp.array(5)) == 10.0

try:
    jax.jit(conditional_func_py)(jnp.array(5.0))
except jax.errors.ConcretizationTypeError as e:
    print("类型:", type(e).__name__)
    print("继承链:", [c.__name__ for c in type(e).__mro__[:3]])

def loop_n(x, n):
    for _ in range(n):
        x = x + 1
    return x

expect_error(lambda: jax.jit(loop_n)(1.0, 3), jax.errors.TracerIntegerConversionError)

# 验证继承关系：只有 Bool 是 ConcretizationTypeError 的子类，四者的共同父类是 JAXTypeError
E = jax.errors
assert issubclass(E.TracerBoolConversionError, E.ConcretizationTypeError)
assert not issubclass(E.TracerIntegerConversionError, E.ConcretizationTypeError)
assert not issubclass(E.TracerArrayConversionError, E.ConcretizationTypeError)
assert all(issubclass(c, E.JAXTypeError) and issubclass(c, TypeError) for c in
           (E.ConcretizationTypeError, E.TracerBoolConversionError,
            E.TracerIntegerConversionError, E.TracerArrayConversionError))
print("继承关系已验证：统一捕获请用 jax.errors.JAXTypeError")

body = lambda i, s: s + i
print("静态上界 → ", "scan" if "scan[" in str(jax.make_jaxpr(lambda s: lax.fori_loop(0, 3, body, s))(0.0)) else "while")
print("追踪上界 → ", "while" if "while[" in str(jax.make_jaxpr(lambda n: lax.fori_loop(0, n, body, 0.0))(3)) else "scan")

# 静态上界：可以求导。d/dx x^3 在 x=2 处 = 12
g_static = jax.grad(lambda x: lax.fori_loop(0, 3, lambda i, s: s * x, 1.0))(2.0)
print("grad(静态 fori):", g_static)
assert g_static == 12.0

# while_loop：反向求导报错
expect_error(lambda: jax.grad(lambda x: lax.while_loop(lambda s: s < 10, lambda s: s * x, 1.0))(2.0), ValueError)

branches = [lambda x: x, lambda x: x * 10, lambda x: -x]
sw = jax.jit(lambda i, x: lax.switch(i, branches, x))
print("switch(1):", sw(1, 3.0), "  switch(99) 被夹到最后一个分支:", sw(99, 3.0))

# while_loop：从 1 开始不断翻倍，直到 >= 10，同时统计步数
final, steps = lax.while_loop(lambda s: s[0] < 10, lambda s: (s[0] * 2, s[1] + 1), (1, 0))
print("while_loop:", final, "步数", steps)

# scan：累加和
carry, ys = lax.scan(lambda c, x: (c + x, c + x), 0.0, jnp.arange(1.0, 5.0))
print("scan: 最终 carry =", carry, " 每步输出 =", ys)

# 分支输出类型不一致会报错
expect_error(lambda: lax.cond(True, lambda x: x, lambda x: jnp.int32(1), 1.0), TypeError)

relu2 = jax.jit(lambda x: jnp.where(x > 0, x * 2, x))
print("jnp.where 逐元素:", relu2(jnp.array([-1.0, 2.0])))
assert jnp.array_equal(relu2(jnp.array([-1.0, 2.0])), jnp.array([-1.0, 4.0]))

# vmap 下带批量谓词的 lax.cond 会变成 select_n：两个分支都会计算
jaxpr_v = str(jax.make_jaxpr(jax.vmap(lambda x: lax.cond(x > 0, lambda v: v * 2, lambda v: -v, x)))(jnp.ones(3)))
print("vmap(cond) 的 Jaxpr 中有 select_n:", "select_n" in jaxpr_v, "| 还有 cond:", "cond[" in jaxpr_v)
assert "select_n" in jaxpr_v and "cond[" not in jaxpr_v

def unrolled(x, n):
    for _ in range(n):
        x = jnp.sin(x) + 1
    return x

def rolled(x, n):
    return lax.fori_loop(0, n, lambda i, x: jnp.sin(x) + 1, x)

for n in [10, 100, 1000]:
    t0 = time.perf_counter(); jax.jit(unrolled, static_argnums=1).lower(1.0, n).compile(); t_u = time.perf_counter() - t0
    t0 = time.perf_counter(); jax.jit(rolled, static_argnums=1).lower(1.0, n).compile(); t_r = time.perf_counter() - t0
    print(f"n={n:5d}  展开: {t_u*1e3:7.1f} ms   fori_loop: {t_r*1e3:6.1f} ms")

def conditional_double(x, threshold):
    if x > threshold:
        return x * 2
    else:
        return x

# 不标记静态参数：x 和 threshold 都被追踪，if 无法得到具体的 True/False
jitted_conditional_double = jax.jit(conditional_double)
expect_error(lambda: jitted_conditional_double(jnp.array(5.0), threshold=0.0), jax.errors.ConcretizationTypeError)

# ❌ 只把 threshold 标记为静态：并不能分别输出 10.0 和 -2.0，而是报错
print("static_argnums=(1,)      →", end=" ")
expect_error(lambda: jax.jit(conditional_double, static_argnums=(1,))(jnp.array(5.0), threshold=0.0),
             jax.errors.ConcretizationTypeError)
print("static_argnames=threshold →", end=" ")
expect_error(lambda: jax.jit(conditional_double, static_argnames=("threshold",))(jnp.array(-2.0), threshold=0.0),
             jax.errors.ConcretizationTypeError)

# ✅ 修法 1：条件依赖 x 的值 → lax.cond 或 jnp.where（threshold 可以是追踪值，换值不重新编译）
@jax.jit
def conditional_double_cond(x, threshold):
    return lax.cond(x > threshold, lambda v: v * 2, lambda v: v, x)

@jax.jit
def conditional_double_where(x, threshold):
    return jnp.where(x > threshold, x * 2, x)

for fn in (conditional_double_cond, conditional_double_where):
    assert fn(jnp.array(5.0), 0.0) == 10.0
    assert fn(jnp.array(-2.0), 0.0) == -2.0
print("修法 1 ✅ lax.cond:", conditional_double_cond(jnp.array(5.0), 0.0), conditional_double_cond(jnp.array(-2.0), 0.0))

# ✅ 修法 2：只有当 if 的条件完全只依赖静态参数时，static_argnames 才真正有效
@partial(jax.jit, static_argnames=["double"])
def maybe_double(x, double):
    if double:            # double 是 Python bool，追踪时已知
        return x * 2
    return x

assert maybe_double(jnp.array(5.0), double=True) == 10.0
assert maybe_double(jnp.array(5.0), double=False) == 5.0
print("修法 2 ✅ 静态开关:", maybe_double(jnp.array(5.0), double=True), maybe_double(jnp.array(5.0), double=False))

def conditional_double_static(x, threshold):
    return jnp.where(x > threshold, x * 2, x)

jitted_func = jax.jit(conditional_double_static, static_argnames=("threshold",))
x5, x1 = jnp.array(5.0), jnp.array(1.0)      # 先建好输入
with count_compiles() as log:
    print("第 1 次调用 (x=5, threshold=0.0)", flush=True);  r_a = jitted_func(x5, threshold=0.0)    # 编译
    print("第 2 次调用 (x=1, threshold=0.0)", flush=True);  r_b = jitted_func(x1, threshold=0.0)    # 复用
    print("第 3 次调用 (x=5, threshold=10.0)", flush=True); r_c = jitted_func(x5, threshold=10.0)   # 重新编译
    print("第 4 次调用 (x=1, threshold=10.0)", flush=True); r_d = jitted_func(x1, threshold=10.0)   # 复用
assert log.count("conditional_double_static") == 2
assert [float(r) for r in (r_a, r_b, r_c, r_d)] == [10.0, 2.0, 5.0, 1.0]
print("✅ 4 次调用共编译", log.count("conditional_double_static"), "次")

from dataclasses import dataclass

# 1) 数组、list 不能作为静态参数；tuple 可以
expect_error(lambda: jax.jit(lambda x, t: x * t, static_argnums=1)(1.0, jnp.array(2.0)), ValueError)
expect_error(lambda: jax.jit(lambda x, shape: x.reshape(shape), static_argnums=1)(jnp.arange(6.0), [2, 3]), ValueError)
reshaped = jax.jit(lambda x, shape: x.reshape(shape), static_argnums=1)(jnp.arange(6.0), (2, 3))
assert reshaped.shape == (2, 3)

# 2) 自定义类作静态参数：按身份哈希 vs 按值哈希
class CfgById:                     # 默认 __hash__/__eq__：按对象身份
    def __init__(self, k):
        self.k = k

@dataclass(frozen=True)
class CfgByValue:                  # frozen dataclass：按字段值哈希
    k: int

def scale_by_id(x, cfg):
    return x * cfg.k

def scale_by_value(x, cfg):
    return x * cfg.k

xc = jnp.ones(3)
j_id = jax.jit(scale_by_id, static_argnums=1)
j_val = jax.jit(scale_by_value, static_argnums=1)
with count_compiles(echo=False) as log:
    j_id(xc, CfgById(2)); j_id(xc, CfgById(2))          # 内容相同但是两个对象
    j_val(xc, CfgByValue(2)); j_val(xc, CfgByValue(2))  # 值相等 → 命中缓存
assert log.count("scale_by_id") == 2
assert log.count("scale_by_value") == 1
print("按身份哈希的配置类编译", log.count("scale_by_id"), "次；frozen dataclass 编译", log.count("scale_by_value"), "次")

# 4) 形状信息是静态的
@jax.jit
def sum_rows_or_all(x):
    return x.sum() if x.ndim == 1 else x.sum(axis=0)

assert sum_rows_or_all(jnp.ones(3)) == 3.0
assert jnp.array_equal(sum_rows_or_all(jnp.ones((2, 3))), jnp.full(3, 2.0))
print("基于 ndim 的 Python if ✅", sum_rows_or_all(jnp.ones(3)), sum_rows_or_all(jnp.ones((2, 3))))

@jax.jit
def condition_on_value(x):
    if x > 0:
        return x * 2
    else:
        return x / 2

expect_error(lambda: condition_on_value(jnp.array(5.0)), jax.errors.ConcretizationTypeError)

# 解决：lax.cond(谓词, 真分支, 假分支, 操作数)
@jax.jit
def condition_on_value_lax(x):
    def true_fun(operand):
        return operand * 2

    def false_fun(operand):
        return operand / 2

    return lax.cond(x > 0, true_fun, false_fun, x)

print(condition_on_value_lax(jnp.array(5.0)))    # 10.0
print(condition_on_value_lax(jnp.array(-4.0)))   # -2.0
assert condition_on_value_lax(jnp.array(5.0)) == 10.0 and condition_on_value_lax(jnp.array(-4.0)) == -2.0

def variable_loop(x, n):
    total = x
    for i in range(n):
        total = total + i
    return total

# n 是追踪值（JAX 数组），range(n) 无法确定迭代次数
expect_error(lambda: jax.jit(variable_loop)(10.0, jnp.array(3)), jax.errors.TracerIntegerConversionError)

# 解决 1：lax.fori_loop（这里配合 static_argnums 把 n 设为静态）
def fixed_loop_fori(x, n_static):
    def body_fun(i, current_total):
        return current_total + i
    return lax.fori_loop(0, n_static, body_fun, x)

jitted_fori = jax.jit(fixed_loop_fori, static_argnums=(1,))
print(jitted_fori(jnp.array(10.0), 3))   # 13.0 = 10 + 0 + 1 + 2
assert jitted_fori(jnp.array(10.0), 3) == 13.0

# 解决 2：lax.scan（带状态的循环，这里求累加和）
@jax.jit
def cumulative_sum_scan(xs):
    def scan_op(carry, x):
        new_carry = carry + x
        return new_carry, new_carry   # (下一步的 carry, 本步输出)
    _, ys = lax.scan(scan_op, 0.0, xs)
    return ys

print(cumulative_sum_scan(jnp.array([1.0, 2.0, 3.0, 4.0])))   # [ 1.  3.  6. 10.]
assert jnp.array_equal(cumulative_sum_scan(jnp.array([1.0, 2.0, 3.0, 4.0])), jnp.array([1.0, 3.0, 6.0, 10.0]))

# ✅ n 作为追踪值传入：结果正确，且不同的 n 共用一份编译结果
def loop_add(x, n):
    return lax.fori_loop(0, n, lambda i, t: t + i, x)

loop_add_jit = jax.jit(loop_add)
x10 = jnp.array(10.0)
with count_compiles() as log:
    results = [float(loop_add_jit(x10, n)) for n in (3, 5, 100)]
print("n=3, 5, 100 →", results)
assert results == [13.0, 20.0, 10.0 + sum(range(100))]
assert log.count("loop_add") == 1
print("✅ 三个不同的 n 只编译了", log.count("loop_add"), "次")

@jax.jit
def function_with_side_effect(x):
    print(f"TRACE: Running Python code for x={x}")   # 只在追踪时执行
    time.sleep(0.1)                                    # 同样只在追踪时执行
    return x * x

print(function_with_side_effect(jnp.array(3.0)))   # 打印 TRACE（x 是追踪器）
print(function_with_side_effect(jnp.array(4.0)))   # 无 TRACE，也不 sleep
print(function_with_side_effect(jnp.array(5.0)))

@jax.jit
def debug_demo(x):
    print("普通 print（仅追踪时）:", x)
    jax.debug.print("jax.debug.print（每次运行）: x = {x}", x=x)
    return x * x

debug_demo(jnp.array(3.0)).block_until_ready()
debug_demo(jnp.array(4.0)).block_until_ready()

# disable_jit：函数里 x 是具体数值，Python if 也能用
with jax.disable_jit():
    print("disable_jit 下:", condition_on_value(jnp.array(5.0)))

@jax.jit
def process_data(x):
    return jnp.sum(x * 2.0)

inputs = [jnp.ones(s) for s in [(3,), (3,), (4,), (4,)]]   # 先建好输入，只计 process_data 本身
with count_compiles(echo=False) as log:
    for x in inputs:
        t0 = time.perf_counter()
        process_data(x).block_until_ready()
        print(f"shape={x.shape}: {time.perf_counter() - t0:.4f} 秒")
assert log.count("process_data") == 2     # (3,) 一次，(4,) 一次
print("✅ 共编译", log.count("process_data"), "次")

@jax.jit
def masked_mean(x_padded, length):
    # length 是追踪值：用掩码代替切片 x[:length]（切片长度必须是静态的）
    mask = jnp.arange(x_padded.shape[0]) < length
    return jnp.sum(jnp.where(mask, x_padded, 0.0)) / length

def bucket_size(n):
    return max(8, 1 << (n - 1).bit_length())   # 向上取整到 2 的幂，最小 8

def mean_any_length(x):
    n = x.shape[0]
    padded = jnp.pad(x, (0, bucket_size(n) - n))
    return masked_mean(padded, n)

lengths = [3, 5, 7, 9, 12, 16, 20]
data_by_len = [jnp.arange(1.0, n + 1) for n in lengths]    # 先建好输入
with count_compiles(echo=False) as log:
    for x in data_by_len:
        assert jnp.allclose(mean_any_length(x), x.mean())
assert log.count("masked_mean") == 3
print("7 种长度结果全部正确；masked_mean 只编译了", log.count("masked_mean"), "次（桶: 8, 16, 32）")

learning_rate = 0.01

@jax.jit
def update_weights(params, grads):
    return params - learning_rate * grads   # learning_rate 在追踪时被固化

params = jnp.array([1.0, 2.0])
grads = jnp.array([0.5, -0.1])

r1 = update_weights(params, grads)
print(r1)                              # [0.995 2.001]
learning_rate = 1000.0                 # 修改全局变量
r2 = update_weights(params, grads)
print(r2)                              # 仍是 [0.995 2.001]：用的是旧值 0.01
assert np.allclose(r1, [0.995, 2.001]) and np.allclose(r2, r1)