JAX 提前编译（AOT）实战指南：拆解 `jax.jit` 背后的 Trace、Lower、Compile 全流程

JAX 提前编译（AOT）实战指南：拆解jax.jit背后的 Trace、Lower、Compile 全流程

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax

JAX 的jax.jit默认在函数被调用时才完成编译，即"即时编译"（Just-In-Time）。本文以 docs/201/aot.md 为主线，讲解 JAX 的 AOT（Ahead-of-Time，提前编译）API：如何把编译拆成trace → lower → compile → execute四个显式步骤，在真正运行之前完成编译、查询 FLOP 估算与内存占用，并掌握静态参数、eval_shape以及"已编译函数不可再被变换"等边界行为。读完本文，你将能够用 AOT API 精确控制编译时机，为服务部署、资源规划与性能调试提供扎实的实战能力。

为什么需要 AOT：从 JIT 到提前编译

jax.jit返回一个"被包装"的函数：当它被调用时，JAX 会当场编译计算并把它运行到加速器（或 CPU）上。正如 JIT 这个缩写所示，所有编译都发生在"为了执行而进行"的那一刻。

但有些场景要求把编译提前到执行之前，或者希望精确控制编译流程中各个阶段的发生时机：

• 部署阶段希望预编译一份可执行程序，运行时不再付出编译开销；
• 希望分别查看/检查trace 产物（jaxpr）、lower 产物（StableHLO）与编译产物（可执行文件）；
• 希望在真正运行之前就知道大概的计算量（FLOPs）与内存占用，用于容量规划。

为此，JAX 的 AOT API 对编译管线中的每一步都提供了直接控制。

jax.jit背后的四个编译阶段

假设F是某个 Python 可调用对象，f = jax.jit(F)。当f(x, y)被调用（其中x、y是数组）时，JAX 依序执行以下四步：

1. Stage out（分离出计算）：基于x、y的类型属性（通常是 shape 与 dtype）推断输入类型，把F的特化版本"分离"成 JAX 的内部中间表示。这一步由 JAX 的tracing（追踪）机制完成，产物是一个jaxpr——JAX 中间语言中的函数（关于 tracing 的概念可参考 tracing）。
2. Lower（降低）：把特化后的计算降低到 XLA 编译器的输入语言StableHLO。
3. Compile（编译）：编译降低后的 HLO 程序，为目标设备（CPU、GPU 或 TPU）生成优化后的可执行文件。
4. Execute（执行）：以x、y为参数执行编译好的可执行文件。

JAX 的 AOT API 恰好让你能单独驱动这四步中的前三步，并在任意阶段停下来检查中间产物。

第一个 AOT 示例：显式走完编译管线

以文档中的经典示例为例：f(x, y) = 2 * x + y，我们用 AOT API 逐步走完整个流程：

>>> import jax >>> import jax.numpy as jnp >>> import numpy as np >>> def f(x, y): return 2 * x + y >>> x, y = 3, 4 >>> traced = jax.jit(f).trace(x, y) >>> # Print the specialized, staged-out representation (as Jaxpr IR) >>> print(traced.jaxpr) { lambda ; a:i32[] b:i32[]. let c:i32[] = mul 2:i32[] a d:i32[] = add c b in (d,) } >>> lowered = traced.lower() >>> # Print lowered HLO >>> print(lowered.as_text()) module @jit_f attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { func.func public @main(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32> {jax.result_info = "result"}) { %c = stablehlo.constant dense<2> : tensor<i32> %0 = stablehlo.multiply %c, %arg0 : tensor<i32> %1 = stablehlo.add %0, %arg1 : tensor<i32> return %1 : tensor<i32> } } >>> compiled = lowered.compile() >>> # Query for cost analysis, print FLOP estimate >>> compiled.cost_analysis()['flops'] 2.0 >>> # Execute the compiled function! >>> compiled(x, y) Array(10, dtype=int32, weak_type=True)

这个例子清晰展示了每个阶段的产物形态：

• trace产物（jaxpr）：mul、add两个等式（eqn）构成的特化计算图，输入类型是i32[]标量。这里2被内联为常量2:i32[]；
• lower产物（StableHLO 文本）：stablehlo.constant、stablehlo.multiply、stablehlo.add，模块头还带有mhlo.num_partitions与mhlo.num_replicas属性；
• compile产物（可执行对象）：可以直接以数组调用，得到结果Array(10, ...)。

从源码结构看，jax.jit(f)返回的对象实现了 jax/_src/stages.py 中定义的Wrappedprotocol：它既是可调用对象（调用即触发 JIT 全流程），又显式暴露trace(...)与lower(...)方法，其中lower(*args, **kwargs)就是trace(*args, **kwargs).lower()的快捷方式。三个阶段的产物分别对应Traced、Lowered、Compiled三个类（公开 API 见 jax/stages.py），它们都继承自带args_info/in_tree/in_avals/donate_argnums属性的Stage基类（jax/_src/stages.py）。

编译期的 XLA 标志：compiler_options

compile步骤接受的compiler_options字典与jax.jit本身一致，用于按次编译设置 XLA 标志。从 jax/_src/api.py 的jit签名可以看到，compiler_options: dict[str, Any] | None = None与in_shardings、out_shardings、static_argnums、donate_argnums、keep_unused、device、backend、inline等参数并列；在 jax/_src/pjit.py 中，这些键值对被展开为tuple并一路传入 XLA 编译后端。想了解 XLA 标志的完整用法（包括全局 vs 按函数设置），可参见 controlling-xla。

运行前查询：FLOP 估算与内存分析

cost_analysis之外，编译后的可执行对象还能在真正运行之前报告内存占用明细，这对判断程序能否装进设备显存非常有用：

>>> stats = compiled.memory_analysis() >>> stats.argument_size_in_bytes, stats.output_size_in_bytes, stats.temp_size_in_bytes (8, 4, 0)

上述例子中，两个int32标量参数合计 8 字节，输出 4 字节，临时内存为 0。这些方法由 jax/_src/stages.py 中Compiled类的cost_analysis()/memory_analysis()实现，它们封装了对底层Executable的查询，并在底层抛出NotImplementedError时返回None（即功能不可用）。Lowered对象也提供cost_analysis()（jax/_src/stages.py），用于在编译之前估算未经编译器优化时的执行代价——注意文档注释特别提醒：该估算发生在编译器优化之前，优化可能大幅改变实际代价，优化后的代价请以Compiled.cost_analysis()为准。

eval_shape：只做 shape/dtype 推断

如果只需要函数的输出类型，既不想 lowering，也不想 compile 或执行，那么eval_shape只运行特化（trace）这一步：

>>> jax.jit(f).eval_shape(jax.ShapeDtypeStruct((), 'int32'), ... jax.ShapeDtypeStruct((), 'int32')) ShapeDtypeStruct(shape=(), dtype=int32)

对于未 jit 的函数，等价功能由jax.eval_shape提供。从源码看，jax.eval_shape的实现非常直观（jax/_src/api.py）：它通过 JAX 的抽象解释机制只做形状推断、不执行任何 FLOPs——对PjitFunction直接调用fun.trace(*args, **kwargs).out_info，其余情况则等价于jit(fun).trace(*args, **kwargs).out_info。其 docstring 明确给出了等价语义：

def eval_shape(fun, *args, **kwargs): out = fun(*args, **kwargs) return jax.tree_util.tree_map(jax.ShapeDtypeStruct.like, out)

同时它也会像真实求值一样抛出 shape 错误，可用于提前捕获形状不匹配的问题。eval_shape与trace返回对象的out_info属性（一棵叶子为ShapeDtypeStruct的 pytree）在 jax/_src/stages.py 中构建。

用ShapeDtypeStruct抽象化参数：不必提供真实数组

jit的所有可选参数——例如static_argnums——在对应的 tracing、lowering、compilation 与执行中都会被尊重。同时，trace的参数不一定非得是真实数组：只要对象带有shape与dtype属性即可：

>>> i32_scalar = jax.ShapeDtypeStruct((), jnp.dtype('int32')) >>> jax.jit(f).trace(i32_scalar, i32_scalar).lower().compile()(x, y) Array(10, dtype=int32)

更一般地，trace只要求其参数在结构上提供 JAX 特化与 lowering 所需的全部信息：对普通数组参数而言，就是shape和dtype字段；而对静态参数而言，JAX 需要的是实际值（详见下一节）。这为"无数据预编译"打开了空间：你可以用ShapeDtypeStruct占位完成编译，再在运行时传入真实数据。

类型不匹配与已编译函数的约束

用与 tracing 时不兼容的参数调用 AOT 编译好的函数会直接报错：

>>> x_1d = y_1d = jnp.arange(3) >>> jax.jit(f).trace(i32_scalar, i32_scalar).lower().compile()(x_1d, y_1d) # doctest: +IGNORE_EXCEPTION_DETAIL ... Traceback (most recent call last): TypeError: Argument types differ from the types for which this computation was compiled. The mismatches are: Argument 'x' compiled with int32[] and called with int32[3] Argument 'y' compiled with int32[] and called with int32[3] >>> x_f = y_f = jnp.float32(72.) >>> jax.jit(f).trace(i32_scalar, i32_scalar).lower().compile()(x_f, y_f) # doctest: +IGNORE_EXCEPTION_DETAIL ... Traceback (most recent call last): TypeError: Argument types differ from the types for which this computation was compiled. The mismatches are: Argument 'x' compiled with int32[] and called with float32[] Argument 'y' compiled with int32[] and called with float32[]

从源码看，Compiled.__call__在调用底层可执行对象前会做两道检查（jax/_src/stages.py）：先比较调用时输入 pytree 结构与编译时的in_tree是否一致，不一致时给出逐条 mismatch 说明（Function compiled with input pytree does not match the input pytree it was called with...）；再在处于变换上下文中时检查是否存在Tracer类型的参数。

静态参数下的 tracing

静态参数最直观地体现了jax.jit的选项、trace的参数、以及最终编译函数所需参数三者之间的相互作用：

>>> lowered_with_x = jax.jit(f, static_argnums=0).trace(7, 8).lower() >>> # Lowered HLO, specialized to the *value* of the first argument (7) >>> print(lowered_with_x.as_text()) module @jit_f attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { func.func public @main(%arg0: tensor<i32>) -> (tensor<i32> {jax.result_info = "result"}) { %c = stablehlo.constant dense<14> : tensor<i32> %0 = stablehlo.add %c, %arg0 : tensor<i32> return %0 : tensor<i32> } } >>> lowered_with_x.compile()(5) Array(19, dtype=int32, weak_type=True)

注意：trace这里照常接收两个参数，但编译后的函数只接收剩下的非静态第二个参数。静态的第一个参数（值 7）在 lowering 时被当作常量，参与常量折叠——它乘以 2 被化简为常量 14，于是 HLO 中只剩下stablehlo.constant dense<14>与stablehlo.add。

虽然trace的第二个参数可以被"空心"的 shape/dtype 结构替换，静态第一个参数必须是具体值，否则 tracing 直接报错：

>>> jax.jit(f, static_argnums=0).trace(i32_scalar, i32_scalar) # doctest: +SKIP Traceback (most recent call last): TypeError: unsupported operand type(s) for *: 'int' and 'ShapeDtypeStruct' >>> jax.jit(f, static_argnums=0).trace(10, i32_scalar).lower().compile()(5) Array(25, dtype=int32)

关于静态参数的完整语义（static_argnums/static_argnames必须可哈希、作为编译缓存键的一部分、inspect.signature的匹配规则等），可查看 jax/_src/api.py 中jit的参数文档，以及 jit 教程。

另外需要注意：trace与lower的产物不能直接序列化后跨进程使用——Lowered.as_text()的 docstring 也明确说明其文本输出"不需要是合法且可靠的序列化"（jax/_src/stages.py）。如果需要可靠的、可移植的序列化，请使用 export 中介绍的 API。

已编译函数不能再被变换（transformed）

编译产物针对一组特定的参数 JAX 类型做了特化（例如特定 shape 与 dtype 的数组）。从 JAX 内部视角看，jax.vmap之类的变换会改变函数的类型签名，使之为编译时的签名所不容。因此 JAX 的策略是：禁止已编译函数参与任何变换。示例：

>>> def g(x): ... assert x.shape == (3, 2) ... return x @ jnp.ones(2) >>> def make_z(*shape): ... return jnp.arange(np.prod(shape)).reshape(shape) >>> z, zs = make_z(3, 2), make_z(4, 3, 2) >>> g_jit = jax.jit(g) >>> g_aot = jax.jit(g).trace(z).lower().compile() >>> jax.vmap(g_jit)(zs) Array([[ 1., 5., 9.], [13., 17., 21.], [25., 29., 33.], [37., 41., 45.]], dtype=float32) >>> jax.vmap(g_aot)(zs) # doctest: +SKIP Traceback (most recent call last): TypeError: Cannot apply JAX transformations to a function lowered and compiled for a particular signature. Detected argument of Tracer type <class 'jax._src.interpreters.batching.BatchTracer'>

g_jit（即时编译版本）可以被vmap批量映射，因为每次调用时它会按新的类型签名重新编译；而g_aot是"为某一签名"定制的产物，vmap传入的BatchTracer无法匹配。同样的错误在g_aot参与自动微分（如jax.grad）时也会出现。为了保持一致，即使jit并不实质改变参数类型签名，对g_aot套用jax.jit同样被禁止。

这个策略在源码中有明确的实现：Compiled.call在非干净 trace 状态下扫描扁平化参数，一旦发现core.Tracer实例即抛出上述TypeError（jax/_src/stages.py）。因此 AOT 编译产物适合作为"最终形态"调用，一切需要变换的逻辑应放在编译之前完成。

调试信息与分析：可用性与不可靠性并存的注意事项

除了核心 AOT 功能（分离且显式的 lowering、编译、执行），各阶段对象还提供一批辅助调试、获取编译器反馈的接口：

• Lowered对象：as_text(dialect=None, *, debug_info=False)输出 lowering 的文本表示，debug_info=True时附带源码位置等调试信息；compiler_ir(dialect=None)返回底层编译器 IR 对象（不可用时返回None）；cost_analysis()返回代价估算（jax/_src/stages.py）；
• Compiled对象：as_text()、cost_analysis()、memory_analysis()、runtime_executable()，以及in_avals/out_info/input_shardings/output_shardings/input_formats/output_formats等属性（jax/_src/stages.py）。

所有这些方法都只是供人工检查与调试的辅助手段，不是可靠的编程接口：它们的可用性与输出随编译器、平台、运行时而变化。由此带来两个重要告诫：

1. 功能不可用：如果 JAX 当前后端不提供某项能力，对应方法返回平凡值（且表现为False类）。例如底层编译器不提供代价分析时，compiled.cost_analysis()返回None。
2. 功能可用但无一致性保证：返回值的类型、结构或数值都不保证在 JAX 配置、后端/平台、版本之间，甚至在同方法的多次调用之间保持一致。今天compiled.cost_analysis()的输出，明天未必相同。

不确定时，请查阅 jax.stages 的包级 API 文档。在 jax/_src/stages.py 中Compiled的 docstring 也明确写道：这些方法的输出"可能是任意的简单数据结构（如嵌套 dict、list、tuple 加数值叶子）……结构可能在不同版本、甚至不同调用之间不一致"。

下一步

本文覆盖了jax.jit底层的各阶段；性能文档的主线接下来是 control-flow——讲解如何在编译代码内部表达条件分支与循环。若要把 lowering 或 compiled 产物序列化到另一进程使用，请参见 export。

附注：本文全部代码示例与行为描述均可在仓库源码中验证——阶段对象定义见 jax/_src/stages.py，jit与eval_shape的签名与实现见 jax/_src/api.py，Wrapped/JitWrapped的trace、lower协议见 jax/_src/pjit.py，公开类型导出见 jax/stages.py。

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax