jax.stages 模块#

编译执行过程各个阶段的接口。

用于即时编译执行的 JAX 变换,例如 jax.jitjax.pmap,也支持一种常见的提前显式降低和编译方法。此模块定义了表示此过程各个阶段的类型。

更多信息,请参阅 AOT 演练

#

class jax.stages.Wrapped(args, **kwargs)[source]#

一个准备被追踪、降低和编译的函数。

此协议反映了诸如 jax.jit 之类函数的输出。调用它会导致 JIT(即时)降低、编译和执行。它也可以在编译之前显式降低,并在执行之前编译结果。

__call__(*args, **kwargs)[source]#

执行包装后的函数,根据需要进行降级和编译。

lower(*args, **kwargs)[source]#

显式地为给定的参数降低此函数。

降低后的函数会从 Python 中分阶段提取,并转换为编译器的输入语言,可能以后端相关的方式进行。它已准备好进行编译,但尚未编译。

返回值:

一个表示降低操作的 Lowered 实例。

返回类型:

Lowered

trace(*args, **kwargs)[source]#

显式地为给定的参数跟踪此函数。

跟踪后的函数会从 Python 中分阶段提取,并转换为 jaxpr。它已准备好进行降级,但尚未降级。

返回值:

一个表示跟踪操作的 Traced 实例。

返回类型:

跟踪

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)

as_text(dialect=None)[source]#

此降级的可读文本表示形式。

用于可视化和调试目的。这不必是有效的或可靠的序列化。它直接转发给外部调用者。

参数:

dialect (str | None | None) – 指定降低方言(例如“stablehlo”)的可选字符串

返回类型:

str

compile(compiler_options=None)[source]#

编译,返回相应的 Compiled 实例。

参数:

compiler_options (CompilerOptions | None | None)

返回类型:

Compiled

compiler_ir(dialect=None)[source]#

此降级的任意对象表示形式。

用于调试目的。这不是有效的或可靠的序列化。输出没有保证跨调用的持续性。

如果不可用,则返回 None,例如基于后端、编译器或运行时。

参数:

dialect (str | None | None) – 指定降低方言(例如“stablehlo”)的可选字符串

返回类型:

Any | None

cost_analysis()[source]#

执行成本估算的摘要。

用于可视化和调试目的。由此输出的对象是一些简单的易于打印或序列化的数据结构(例如,嵌套字典、列表和元组,叶子为数值)。但是,其结构可以是任意的:它可能在 JAX 和 jaxlib 的不同版本之间,甚至在不同的调用之间不一致。

如果不可用,则返回 None,例如基于后端、编译器或运行时。

返回类型:

Any | None

property in_tree: tree_util.PyTreeDef[source]#

(位置参数、关键字参数)对的树结构。

class jax.stages.Compiled(executable, args_info, out_tree, no_kwargs=False)[source]#

专门针对类型/值的函数的编译表示形式。

编译后的计算与可执行文件以及执行它所需的其他信息相关联。它还为跨 JAX 的各种编译路径和后端查询编译计算属性提供了一个通用 API。

参数:
  • args_info (Any)

  • out_tree (tree_util.PyTreeDef)

__call__(*args, **kwargs)[source]#

将自身作为函数调用。

as_text()[source]#

此可执行文件的可读文本表示形式。

用于可视化和调试目的。这不是有效的或可靠的序列化。

如果不可用,则返回 None,例如基于后端、编译器或运行时。

返回类型:

str | None

cost_analysis()[source]#

执行成本估算的摘要。

用于可视化和调试目的。由此输出的对象是一些简单的易于打印或序列化的数据结构(例如,嵌套字典、列表和元组,叶子为数值)。但是,其结构可以是任意的:它可能在 JAX 和 jaxlib 的不同版本之间,甚至在不同的调用之间不一致。

如果不可用,则返回 None,例如基于后端、编译器或运行时。

返回类型:

Any | None

property in_tree: tree_util.PyTreeDef[source]#

(位置参数、关键字参数)对的树结构。

memory_analysis()[source]#

估计内存需求的摘要。

用于可视化和调试目的。由此输出的对象是一些简单的易于打印或序列化的数据结构(例如,嵌套字典、列表和元组,叶子为数值)。但是,其结构可以是任意的:它可能在 JAX 和 jaxlib 的不同版本之间,甚至在不同的调用之间不一致。

如果不可用,则返回 None,例如基于后端、编译器或运行时。

返回类型:

Any | None

runtime_executable()[source]#

此可执行文件的任意对象表示形式。

用于调试目的。这不是有效的或可靠的序列化。输出没有保证跨调用的持续性。

如果不可用,则返回 None,例如基于后端、编译器或运行时。

返回类型:

Any | None