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

JAX 提前编译AOT实战指南拆解jax.jit背后的 Trace、Lower、Compile 全流程【免费下载链接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jaxJAX 的jax.jit默认在函数被调用时才完成编译即即时编译Just-In-Time。本文以 docs/201/aot.md 为主线讲解 JAX 的 AOTAhead-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 依序执行以下四步Stage out分离出计算基于x、y的类型属性通常是 shape 与 dtype推断输入类型把F的特化版本分离成 JAX 的内部中间表示。这一步由 JAX 的tracing追踪机制完成产物是一个jaxpr——JAX 中间语言中的函数关于 tracing 的概念可参考 tracing。Lower降低把特化后的计算降低到 XLA 编译器的输入语言StableHLO。Compile编译编译降低后的 HLO 程序为目标设备CPU、GPU 或 TPU生成优化后的可执行文件。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: tensori32, %arg1: tensori32) - (tensori32 {jax.result_info result}) { %c stablehlo.constant dense2 : tensori32 %0 stablehlo.multiply %c, %arg0 : tensori32 %1 stablehlo.add %0, %arg1 : tensori32 return %1 : tensori32 } } 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, dtypeint32, weak_typeTrue)这个例子清晰展示了每个阶段的产物形态trace产物jaxprmul、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_optionscompile步骤接受的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(), dtypeint32)对于未 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, dtypeint32)更一般地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_argnums0).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: tensori32) - (tensori32 {jax.result_info result}) { %c stablehlo.constant dense14 : tensori32 %0 stablehlo.add %c, %arg0 : tensori32 return %0 : tensori32 } } lowered_with_x.compile()(5) Array(19, dtypeint32, weak_typeTrue)注意trace这里照常接收两个参数但编译后的函数只接收剩下的非静态第二个参数。静态的第一个参数值 7在 lowering 时被当作常量参与常量折叠——它乘以 2 被化简为常量 14于是 HLO 中只剩下stablehlo.constant dense14与stablehlo.add。虽然trace的第二个参数可以被空心的 shape/dtype 结构替换静态第一个参数必须是具体值否则 tracing 直接报错 jax.jit(f, static_argnums0).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_argnums0).trace(10, i32_scalar).lower().compile()(5) Array(25, dtypeint32)关于静态参数的完整语义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.]], dtypefloat32) 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.BatchTracerg_jit即时编译版本可以被vmap批量映射因为每次调用时它会按新的类型签名重新编译而g_aot是为某一签名定制的产物vmap传入的BatchTracer无法匹配。同样的错误在g_aot参与自动微分如jax.grad时也会出现。为了保持一致即使jit并不实质改变参数类型签名对g_aot套用jax.jit同样被禁止。这个策略在源码中有明确的实现Compiled.call在非干净 trace 状态下扫描扁平化参数一旦发现core.Tracer实例即抛出上述TypeErrorjax/_src/stages.py。因此 AOT 编译产物适合作为最终形态调用一切需要变换的逻辑应放在编译之前完成。调试信息与分析可用性与不可靠性并存的注意事项除了核心 AOT 功能分离且显式的 lowering、编译、执行各阶段对象还提供一批辅助调试、获取编译器反馈的接口Lowered对象as_text(dialectNone, *, debug_infoFalse)输出 lowering 的文本表示debug_infoTrue时附带源码位置等调试信息compiler_ir(dialectNone)返回底层编译器 IR 对象不可用时返回Nonecost_analysis()返回代价估算jax/_src/stages.pyCompiled对象as_text()、cost_analysis()、memory_analysis()、runtime_executable()以及in_avals/out_info/input_shardings/output_shardings/input_formats/output_formats等属性jax/_src/stages.py。所有这些方法都只是供人工检查与调试的辅助手段不是可靠的编程接口它们的可用性与输出随编译器、平台、运行时而变化。由此带来两个重要告诫功能不可用如果 JAX 当前后端不提供某项能力对应方法返回平凡值且表现为False类。例如底层编译器不提供代价分析时compiled.cost_analysis()返回None。功能可用但无一致性保证返回值的类型、结构或数值都不保证在 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.pyjit与eval_shape的签名与实现见 jax/_src/api.pyWrapped/JitWrapped的trace、lower协议见 jax/_src/pjit.py公开类型导出见 jax/stages.py。【免费下载链接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考