JAX 的 jit 编译会将函数抽象追踪（trace time）一次生成静态计算图，后续调用直接执行编译后代码（runtime），导致普通 print 仅在首次追踪时触发；应改用 jax.debug.print 实现每次迭代的运行时输出。
jax 的 `print` 在 `grad` 和 `jit` 下失效的原因与解决方案：jax 的 `jit` 编译会将函数抽象追踪（trace time）一次生成静态计算图，后续调用直接执行编译后代码（runtime），导致普通 `print` 仅在首次追踪时触发；应改用 `jax.debug.print` 实现每次迭代的运行时输出。
在 JAX 中，函数的 JIT 编译并非简单“加速执行”，而是一个两阶段过程：追踪（tracing） 和 运行（execution）。当你首次调用 jax.jit(jax.grad(cost))(params) 时，JAX 会以抽象值（如 ShapedArray）为输入对 cost 进行符号化执行，构建计算图并编译为高效内核——这一阶段称为 trace time。此时所有 Python 副作用（如 print()、logging.info()、全局变量修改）都会被执行，因此你看到第一次 "Evaluating" 输出。
但从第二次调用开始，JAX 直接复用已编译的内核，跳过 Python 层逻辑，进入纯 runtime 阶段：输入真实数值，经底层 XLA 执行梯度计算，不再重新解释 Python 语句。于是 print('Evaluating') 永远不会再次执行——它已被“编译掉”了。
✅ 正确做法是使用 JAX 原生调试工具：jax.debug.print。该函数被专门设计为可追踪的运行时打印原语，其调用会被插入到计算图中，并在每次 runtime 执行时触发（支持格式化、条件打印、设备端日志等）：
def cost(params):
jax.debug.print("Evaluating at params = {x}", x=params) # ✅ 每次迭代都输出
return circuit(params)? 提示：
jax.debug.print默认同步输出（阻塞执行），适合调试；若需异步高性能日志，可搭配jax.debug.callback或jax.experimental.host_callback（注意后者在 pmap/pjit 中行为受限）。
⚠️ 注意事项：
- 不要依赖
print()、logging或任何非 JAX-aware 的副作用进行调试或状态观测； -
jax.debug.print的参数必须是 JAX 兼容数组（如jnp.ndarray），不可传入未转换的 Python 原生类型（如float、list），否则报错； - 在
@jax.jit或jax.grad包裹的函数中，所有逻辑应保持函数式（无状态、无外部依赖），副作用必须显式使用jax.debug.*系列 API； - 若需更细粒度控制（如仅在特定 epoch 打印），可结合
jax.debug.callback+ 条件判断，但需确保回调函数本身是纯且可追踪的。
总结：JAX 的“静默优化”不是 bug，而是其函数式与编译模型的必然体现。拥抱 jax.debug.print 等专用调试接口，是编写可维护、可复现 JAX 程序的关键实践。
