jax.stages 模块#

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

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

更多信息,请参阅 AOT 演练

#

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

一个准备好被跟踪 (traced)、降低 (lowered) 和编译的函数。

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

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

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

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

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

这是 self.trace(*args, **kwargs).lower() 的快捷方式。

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

返回:

一个 Lowered 实例,表示降低。

返回类型:

Lowered

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

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

跟踪的函数从 Python 中 staged out,并转换为 jaxpr。它已准备好降低,但尚未降低。

返回:

一个 Traced 实例,表示跟踪。

返回类型:

Traced

class jax.stages.Traced(jaxpr, args_info, fun_name, out_tree, lower_callable, args_flat=None, arg_names=None, num_consts=0)[source]#

函数针对参数类型和值进行专门化的跟踪形式。

跟踪的计算已准备好进行降低。此类携带跟踪表示以及稍后降低、编译和执行它所需的其余信息。

参数:
lower(*, lowering_platforms=None, _private_parameters=None)[source]#

降低到编译器输入,返回一个 Lowered 实例。

参数:
  • lowering_platforms (tuple[str, ...] | None | None)

  • _private_parameters (mlir.LoweringParameters | None | None)

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

函数针对参数类型和值进行专门化的降低形式。

降低是准备好进行编译的计算。此类携带降低以及稍后编译和执行它所需的其余信息。它还为查询 JAX 各种降低路径(jit()pmap() 等)的降低计算的属性提供了一个通用 API。

参数:
  • lowering (Lowering)

  • args_info (Any)

  • out_tree (tree_util.PyTreeDef)

  • no_kwargs (bool)

as_text(dialect=None, *, debug_info=False)[source]#

此降低的人类可读文本表示形式。

旨在用于可视化和调试目的。这不需要是有效的或可靠的序列化。如果您想要可靠且可移植的序列化,请使用 jax.export

参数:
  • dialect (str | None | None) – 可选字符串,指定降低方言(例如“stablehlo”或“hlo”)。

  • debug_info (bool) – 是否包含调试信息,例如,源位置。

返回类型:

str

compile(compiler_options=None)[source]#

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

参数:

compiler_options (CompilerOptions | None | None)

返回类型:

Compiled

compiler_ir(dialect=None)[source]#

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

旨在用于调试目的。这不是有效且可靠的序列化。输出不能保证在调用之间保持一致。如果您想要可靠且可移植的序列化,请使用 jax.export

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

参数:

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

返回类型:

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