jax.export 模块#

jax.export 是一个用于导出和序列化 JAX 函数以便永久存档的库。

请参阅 导出和序列化 文档。

类#

class jax.export.Exported(fun_name, in_tree, in_avals, out_tree, out_avals, in_shardings_hlo, out_shardings_hlo, nr_devices, platforms, ordered_effects, unordered_effects, disabled_safety_checks, mlir_module_serialized, calling_convention_version, module_kept_var_idx, uses_global_constants, _get_vjp)[source]#

已降级为 StableHLO 的 JAX 函数。

参数:

  • fun_name (str)

  • in_tree (tree_util.PyTreeDef)

  • in_avals (tuple[core.ShapedArray, ...])

  • out_tree (tree_util.PyTreeDef)

  • out_avals (tuple[core.ShapedArray, ...])

  • in_shardings_hlo (tuple[HloSharding | None, ...])

  • out_shardings_hlo (tuple[HloSharding | None, ...])

  • nr_devices (int)

  • platforms (tuple[str, ...])

  • ordered_effects (tuple[effects.Effect, ...])

  • unordered_effects (tuple[effects.Effect, ...])

  • disabled_safety_checks (Sequence[DisabledSafetyCheck])

  • mlir_module_serialized (bytes)

  • calling_convention_version (int)

  • module_kept_var_idx (tuple[int, ...])

  • uses_global_constants (bool)

  • _get_vjp (Callable[[Exported], Exported] | None)

fun_name#

导出的函数的名称，用于错误消息。

类型:

str

in_tree#

描述已降级 JAX 函数的（args, kwargs）元组的 PyTreeDef。实际降级不依赖于 in_tree，但可用于使用相同的参数结构调用导出的函数。

类型:

tree_util.PyTreeDef

in_avals#

输入的抽象值的扁平元组。可能包含形状中的维度表达式。

类型:

tuple[core.ShapedArray, …]

out_tree#

描述已降级 JAX 函数结果的 PyTreeDef。

类型:

tree_util.PyTreeDef

out_avals#

输出抽象值的扁平元组。可能包含形状中的维度表达式，维度变量位于 in_avals 中。

类型:

tuple[core.ShapedArray, …]

in_shardings_hlo#

扁平化的输入分片，一个长度与 in_avals 相同的序列。None 表示未指定分片。请注意，这些不包括分片或分片中使用的实际设备。有关将这些转换为可与 JAX API 结合使用的分片规范的方法，请参见 in_shardings_jax。

类型:

tuple[HloSharding | None, …]

out_shardings_hlo#

扁平化的输出分片，一个长度与 out_avals 相同的序列。None 表示未指定分片。请注意，这些不包括分片或分片中使用的实际设备。有关将这些转换为可与 JAX API 结合使用的分片规范的方法，请参见 out_shardings_jax。

类型:

tuple[HloSharding | None, …]

nr_devices#

模块已降级的设备数量。

类型:

int

platforms#

包含函数应导出的平台的元组。JAX 的平台集合是开放式的；用户可以添加平台。JAX 内置平台包括：‘tpu’、‘cpu’、‘cuda’、‘rocm’。请参阅 https://jax.net.cn/en/latest/export/export.html#cross-platform-and-multi-platform-export。

类型:

tuple[str, …]

ordered_effects#

序列化模块中存在的有序效应。此项自序列化版本 9 起可用。有关存在有序效应时的调用约定，请参阅 https://jax.net.cn/en/latest/export/export.html#module-calling-convention。

类型:

tuple[effects.Effect, …]

unordered_effects#

序列化模块中存在的无序效应。此项自序列化版本 9 起可用。

类型:

tuple[effects.Effect, …]

mlir_module_serialized#

序列化的已降级 VHLO 模块。

类型:

bytes

calling_convention_version#

已导出模块调用约定的版本号。有关更多版本信息，请参阅 https://jax.net.cn/en/latest/export/export.html#calling-convention-versions。

类型:

int

module_kept_var_idx#

必须传递给模块的 in_avals 中的参数的已排序索引。其他参数已被删除，因为它们未使用。

类型:

tuple[int, …]

uses_global_constants#

是否 mlir_module_serialized 使用形状多态或多平台导出。这可能是因为 in_avals 包含维度变量，或者由于 Exported 模块的内部调用包含维度变量或平台索引参数。此类模块在 XLA 编译之前需要进行形状细化。

类型:

bool

disabled_safety_checks#

已在导出时禁用的安全检查描述符列表。请参阅 DisabledSafetyCheck 的文档字符串。

类型:

Sequence[DisabledSafetyCheck]

_get_vjp#

一个可选函数，它接收当前的导出的函数并返回导出的 VJP 函数。VJP 函数接收一个扁平的参数列表，从原始参数开始，然后是每个原始输出的余切参数。它返回一个包含对应于扁平化原始输入的余切的元组。

类型:

Callable[[Exported], Exported] | None

有关 mlir_module() 方法的调用约定的描述，请参阅 https://jax.net.cn/en/latest/export/export.html#module-calling-convention。

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

从 JAX 程序中调用导出的函数。

参数:

  • args – 要传递给导出的函数的 positional 参数。这应该是一个与函数导出时所用参数具有相同 PyTree 结构的数组的 pytree。

  • kwargs – 要传递给导出的函数的 keyword 参数。

Returns: 结果数组的 pytree，其结构与导出的函数的结果相同。

results of the exported function.

调用支持反向模式 AD 以及导出的所有功能：形状多态、多平台、设备多态。请参阅 [JAX export documentation](https://jax.net.cn/en/latest/export/export.html) 中的示例。

has_vjp()[source]#

返回此 Exported 是否支持 VJP。

返回类型:

bool

in_shardings_jax(mesh)[source]#

创建与 self.in_shardings_hlo 对应的 Shardings。

Exported 对象将 in_shardings_hlo 存储为 HloShardings，它们独立于分片或设备集。此方法构造的 Sharding 可用于 JAX API，例如 jax.jit() 或 jax.device_put()。

示例用法

>>> from jax import export, sharding
>>> # Prepare the exported object:
>>> exp_mesh = sharding.Mesh(jax.devices(), ("a",))
>>> exp = export.export(jax.jit(lambda x: jax.numpy.add(x, x),
...                             in_shardings=sharding.NamedSharding(exp_mesh, sharding.PartitionSpec("a")))
...     )(np.arange(jax.device_count()))
>>> exp.in_shardings_hlo
({devices=[8]<=[8]},)
>>> # Create a mesh for running the exported object
>>> run_mesh = sharding.Mesh(jax.devices()[::-1], ("b",))
>>> # Put the args and kwargs on the appropriate devices
>>> run_arg = jax.device_put(np.arange(jax.device_count()),
...     exp.in_shardings_jax(run_mesh)[0])
>>> res = exp.call(run_arg)
>>> res.addressable_shards
[Shard(device=CpuDevice(id=7), index=(slice(0, 1, None),), replica_id=0, data=[0]),
 Shard(device=CpuDevice(id=6), index=(slice(1, 2, None),), replica_id=0, data=[2]),
 Shard(device=CpuDevice(id=5), index=(slice(2, 3, None),), replica_id=0, data=[4]),
 Shard(device=CpuDevice(id=4), index=(slice(3, 4, None),), replica_id=0, data=[6]),
 Shard(device=CpuDevice(id=3), index=(slice(4, 5, None),), replica_id=0, data=[8]),
 Shard(device=CpuDevice(id=2), index=(slice(5, 6, None),), replica_id=0, data=[10]),
 Shard(device=CpuDevice(id=1), index=(slice(6, 7, None),), replica_id=0, data=[12]),
 Shard(device=CpuDevice(id=0), index=(slice(7, 8, None),), replica_id=0, data=[14])]
参数:

mesh (mesh_lib.Mesh)

返回类型:

Sequence[sharding.Sharding | None]

mlir_module()[source]#

对 mlir_module_serialized 的字符串表示。

返回类型:

str

out_shardings_jax(mesh)[source]#

创建与 self.out_shardings_hlo 对应的 Shardings。

请参阅 in_shardings_jax 的文档。

参数:

mesh (mesh_lib.Mesh)

返回类型:

Sequence[sharding.Sharding | None]

serialize(vjp_order=0)[source]#

序列化一个 Exported 对象。

参数:

vjp_order (int) – 要包含的最大 vjp 阶数。例如，值为 2 表示我们序列化原始函数和两阶 vjp 函数。这应该允许对反序列化的函数进行二阶反向模式微分。即 jax.grad(jax.grad(f))。

返回类型:

bytearray

vjp()[source]#

获取导出的 VJP。

如果 VJP 不可用，则返回 None，这可能发生在 Exported 对象是从没有 VJP 的外部格式加载时。

返回类型:

Exported

class jax.export.DisabledSafetyCheck(_impl)[source]#

在（反）序列化时应跳过的安全检查。

其中大多数检查在序列化时执行，但有些则推迟到反序列化。禁用的检查列表会附加到序列化中，例如作为 jax.export.Exported 的字符串属性序列，或者 tf.XlaCallModuleOp 的字符串属性序列。

使用 jax2tf 时，可以通过传递 TF_XLA_FLAGS=--tf_xla_call_module_disabled_checks=platform 来禁用更多反序列化安全检查。

参数:

_impl (str)

classmethod custom_call(target_name)[source]#

允许序列化一个尚不确定稳定的调用目标。

仅在序列化时生效。 :param target_name: 要允许的自定义调用目标的名称。

参数:

target_name (str)

返回类型:

DisabledSafetyCheck

is_custom_call()[source]#

返回此指令允许的自定义调用目标。

返回类型:

str | None

classmethod platform()[source]#

允许编译平台与导出平台不同。

仅在反序列化时生效。

返回类型:

DisabledSafetyCheck

函数#

export(fun_jit, *[, platforms, ...])

导出 JAX 函数以进行持久化序列化。

deserialize(blob)

反序列化一个 Exported 对象。

minimum_supported_calling_convention_version

int([x]) -> integer int(x, base=10) -> integer

maximum_supported_calling_convention_version

int([x]) -> integer int(x, base=10) -> integer

default_export_platform()

检索默认导出平台。

register_pytree_node_serialization(nodetype, ...)

注册自定义 PyTree 节点以进行序列化和反序列化。

register_namedtuple_serialization(nodetype, ...)

注册一个 namedtuple 以进行序列化和反序列化。

常量#

jax.export.minimum_supported_serialization_version#

最低支持的序列化版本；请参阅 调用约定版本。

jax.export.maximum_supported_serialization_version#

最高支持的序列化版本；请参阅 调用约定版本。