如何完全隔离代码块与JAX tracing，避免辅助计算被追踪

之前针对该问题的表述可能存在模糊，原因是对JAX tracing的工作原理了解有限，无法精准描述需求，本次聚焦的核心疑问即上述问题。

场景背景 #

当前待处理的函数结构如下：

def f(x, y):
   z = h(y)
   return g(x, z)

核心需求 #

在对g(x, z)执行各类JAX变换时，将z作为常量处理。由于直接构造参数z的流程十分繁琐，因此借助辅助函数h，将易于指定的输入y转换为g要求的格式。
期望JAX将h识别为不可追踪的黑盒，最终实现的效果为：

• 针对特定输入y0，执行jit(lambda x: f(x, y0))的效果，完全等价于先用numpy计算得到z0 = h(y0)，再执行jit(lambda x: g(x, z0))
• 对grad或其他各类JAX函数变换，也达到同等效果

实际遇到的问题 #

目前已将h实现为仅使用标准numpy的版本，原本以为这样会自动具备黑盒不被追踪的特性，但实际运行发现：
jit(lambda x: f(x, y0))的编译耗时，明显高于“先计算z0 = h(y0)，再编译jit(lambda x: g(x, z0))”的编译耗时。目前推测编译耗时的差异与JAX对h内部大量循环逻辑的追踪有关，但暂未确认。

不直接改造h为JAX友好实现的原因 #

• h的输入格式不规则，内部包含大量循环与条件判断，输出形状依赖输入取值，改造难度极高
• 该函数执行成本极低，输入为整数类型，完全不需要对其执行求导操作，投入精力改写的性价比极低

不采用“手动预计算z0再做JAX变换”方案的原因 #

该方案扩展性不足：当存在大量结构与f类似、且相互调用形成链式调用的函数，这些函数包含JAX不友好的输入与计算逻辑（但这类输入与计算在JAX变换流程中始终会被视为常量）时，手动提取所有预计算流程的实现难度极高。

最终聚焦的疑问 #

是否存在方法可以控制f的追踪过程，让JAX在追踪f时，直接对z = h(y)执行真实求值（而非逐行追踪h的内部逻辑），之后再继续追踪g(x, z)的计算流程？

内容的提问来源于stack exchange，提问作者DanPuzzuoli