jax.stages
模块#
编译执行过程各个阶段的接口。
用于即时编译执行的 JAX 变换,例如 jax.jit
和 jax.pmap
,也支持一种常见的提前显式降低和编译方法。此模块定义了表示此过程各个阶段的类型。
更多信息,请参阅 AOT 演练。
类#
- class jax.stages.Wrapped(*args, **kwargs)[source]#
一个准备被追踪、降低和编译的函数。
此协议反映了诸如
jax.jit
之类函数的输出。调用它会导致 JIT(即时)降低、编译和执行。它也可以在编译之前显式降低,并在执行之前编译结果。
- class jax.stages.Lowered(lowering, args_info, out_tree, no_kwargs=False)[source]#
专门针对参数类型和值的函数的降级。
降级是一种准备好进行编译的计算。此类将降级操作与稍后编译和执行它所需的其他信息一起携带。它还为跨 JAX 的各种降级路径(
jit()
、pmap()
等)查询降级计算属性提供了一个通用 API。- 参数:
lowering (XlaLowering)
args_info (Any)
out_tree (tree_util.PyTreeDef)
no_kwargs (bool)
- compile(compiler_options=None)[source]#
编译,返回相应的
Compiled
实例。- 参数:
compiler_options (CompilerOptions | None | None)
- 返回类型:
- compiler_ir(dialect=None)[source]#
此降级的任意对象表示形式。
用于调试目的。这不是有效的或可靠的序列化。输出没有保证跨调用的持续性。
如果不可用,则返回
None
,例如基于后端、编译器或运行时。- 参数:
dialect (str | None | None) – 指定降低方言(例如“stablehlo”)的可选字符串
- 返回类型:
Any | None
- class jax.stages.Compiled(executable, args_info, out_tree, no_kwargs=False)[source]#
专门针对类型/值的函数的编译表示形式。
编译后的计算与可执行文件以及执行它所需的其他信息相关联。它还为跨 JAX 的各种编译路径和后端查询编译计算属性提供了一个通用 API。
- 参数:
args_info (Any)
out_tree (tree_util.PyTreeDef)
- as_text()[source]#
此可执行文件的可读文本表示形式。
用于可视化和调试目的。这不是有效的或可靠的序列化。
如果不可用,则返回
None
,例如基于后端、编译器或运行时。- 返回类型:
str | None
- cost_analysis()[source]#
执行成本估算的摘要。
用于可视化和调试目的。由此输出的对象是一些简单的易于打印或序列化的数据结构(例如,嵌套字典、列表和元组,叶子为数值)。但是,其结构可以是任意的:它可能在 JAX 和 jaxlib 的不同版本之间,甚至在不同的调用之间不一致。
如果不可用,则返回
None
,例如基于后端、编译器或运行时。- 返回类型:
Any | None