导出和序列化分阶段计算#
预先降级和编译 (Ahead-of-time lowering and compilation) API 生成的对象可用于调试，或在同一进程中进行编译和执行。有时，您希望将已降级的 JAX 函数序列化，以便在单独的进程中（可能在稍后的时间）进行编译和执行。这使您能够：
在另一个进程或机器上编译并执行函数，而无需访问原始 JAX 程序，也无需重复进行分阶段输出和降级过程（例如在推理系统中）。
在没有访问权限的机器上对函数进行追踪和降级，而在稍后想要在该加速器上编译并执行该函数。
归档 JAX 函数的快照，例如以便日后重现结果。注意：请查看此用例的兼容性保证。
有关更多详细信息，请参阅 jax.export API 参考。
以下是一个示例
>>> import re
>>> import numpy as np
>>> import jax
>>> from jax import export
>>> def f(x): return 2 * x * x
>>> exported: export.Exported = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((), np.float32))
>>> # You can inspect the Exported object
>>> print(re.search(r".*@main.*", exported.mlir_module()).group(0))
func.func public @main(%arg0: tensor<f32> loc("x")) -> (tensor<f32> {jax.result_info = "result"}) {
>>> # And you can serialize the Exported to a bytearray.
>>> serialized: bytearray = exported.serialize()
>>> # The serialized function can later be rehydrated and called from
>>> # another JAX computation, possibly in another process.
>>> rehydrated_exp: export.Exported = export.deserialize(serialized)
>>> def callee(y):
... return 3. * rehydrated_exp.call(y * 4.)
>>> callee(1.)
Array(96., dtype=float32)
序列化分为两个阶段：
导出以生成一个
jax.export.Exported对象，其中包含用于降级函数的 StableHLO，以及从另一个 JAX 函数调用它所需的元数据。我们计划添加从 TensorFlow 生成Exported对象的代码，并使用来自 TensorFlow 和 PyTorch 的Exported对象。使用 flatbuffers 格式将实际数据序列化为字节数组。请参阅 与 TensorFlow 的互操作 (Interoperation with TensorFlow)，了解用于与 TensorFlow 互操作的 TensorFlow 图序列化的替代方案。
支持反向模式自动微分 (AD)#
序列化可以根据需要支持高阶反向模式 AD。这通过序列化原始函数的 jax.vjp() 以及原始函数本身来实现，最高可达用户指定的阶数（默认为 0，意味着重新水合后的函数无法进行微分）。
>>> import jax
>>> from jax import export
>>> from typing import Callable
>>> def f(x): return 7 * x * x * x
>>> # Serialize 3 levels of VJP along with the primal function
>>> blob: bytearray = export.export(jax.jit(f))(1.).serialize(vjp_order=3)
>>> rehydrated_f: Callable = export.deserialize(blob).call
>>> rehydrated_f(0.1) # 7 * 0.1^3
Array(0.007, dtype=float32)
>>> jax.grad(rehydrated_f)(0.1) # 7*3 * 0.1^2
Array(0.21000001, dtype=float32)
>>> jax.grad(jax.grad(rehydrated_f))(0.1) # 7*3*2 * 0.1
Array(4.2, dtype=float32)
>>> jax.grad(jax.grad(jax.grad(rehydrated_f)))(0.1) # 7*3*2
Array(42., dtype=float32)
>>> jax.grad(jax.grad(jax.grad(jax.grad(rehydrated_f))))(0.1)
Traceback (most recent call last):
ValueError: No VJP is available
请注意，VJP 函数是在序列化时延迟计算的（此时 JAX 程序仍然可用）。这意味着它尊重 JAX VJP 的所有特性，例如 jax.custom_vjp() 和 jax.remat()。
请注意，重新水合后的函数不支持任何其他转换，例如前向模式 AD (jvp) 或 jax.vmap()。
兼容性保证#
当您使用 jax.export 模块时，您将获得以下导出兼容性保证：JAX 导出并序列化的工件支持由编译器和 JAX 运行时系统进行 .deserialize 和 .call（即编译和执行），前提是系统版本：
比导出所用的 JAX 版本新不超过 6 个月（我们称 JAX 导出提供 6 个月的后向兼容性）。如果您希望归档导出的工件以便日后编译和执行，这非常有用。
比导出所用的 JAX 版本旧不超过 3 周（我们称 JAX 导出提供 3 周的前向兼容性）。如果您希望使用在导出之前构建和部署的消费者来编译和运行导出的工件（例如，在导出时已部署的推理系统），这非常有用。
（具体的兼容性窗口长度与 JAX 为 jax2tf 承诺的兼容性相同，并基于 TensorFlow 兼容性。“后向兼容性”这一术语是从消费者（例如推理系统）的角度出发的。）
重要的是导出和消费组件的构建时间，而不是导出和编译发生的具体时间。为了减少不兼容的可能性，内部 JAX 用户应尽可能频繁地重建和重新部署消费系统。
外部用户在进行归档导出时应使用最新发布的 jaxlib 版本。
如果您绕过 jax.export API 来获取 StableHLO 代码（例如通过 jax.jit(f).lower(1.).compiler_ir()）并尝试将其用于不同版本的 JAX 或 jaxlib 构建，则兼容性保证不适用。与直接降级不同，jax.export 模块使用 StableHLO 的可移植工件特性来处理 StableHLO 操作集可能的演变。
自定义调用的兼容性保证#
此外，原始 StableHLO 可能包含引用 C++ 函数的自定义调用 (custom calls)。JAX 使用自定义调用来降级少数基元，例如线性代数基元、分片注释 (sharding annotations) 或 Pallas 内核。这些不属于 StableHLO 的兼容性保证范围。这些函数的 C++ 实现很少更改，但它们确实可能更改。
为了确保前向兼容性，当我们更改 JAX 降级规则以使用新的自定义调用目标时，JAX 将在 3 周内避免使用该新目标。要使用最新的降级规则，您可以传递 --jax_export_ignore_forward_compatibility=1 配置标志或设置 JAX_EXPORT_IGNORE_FORWARD_COMPATIBILITY=1 环境变量。
只有一小部分自定义调用被保证稳定并具有兼容性保证（查看列表）。我们不断向允许列表中添加更多的自定义调用目标，并提供后向兼容性测试。如果您尝试序列化调用其他自定义调用目标的代码，将在导出期间收到错误。
如果您想针对特定的自定义调用（例如目标为 my_target）禁用此安全检查，您可以将 export.DisabledSafetyCheck.custom_call("my_target") 添加到 export 方法的 disabled_checks 参数中，如下例所示。
>>> import jax
>>> from jax import export
>>> from jax import lax
>>> from jax._src import core
>>> from jax._src.interpreters import mlir
>>> # Define a new primitive backed by a custom call
>>> new_prim = core.Primitive("new_prim")
>>> _ = new_prim.def_abstract_eval(lambda x: x)
>>> _ = mlir.register_lowering(new_prim, lambda ctx, o: mlir.custom_call("my_new_prim", operands=[o], result_types=[o.type]).results)
>>> print(jax.jit(new_prim.bind).lower(1.).compiler_ir())
module @jit_bind attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<f32>) -> (tensor<f32> {jax.result_info = "result"}) {
%0 = stablehlo.custom_call @my_new_prim(%arg0) {api_version = 2 : i32, backend_config = ""} : (tensor<f32>) -> tensor<f32>
return %0 : tensor<f32>
}
}
>>> # If we try to export, we get an error
>>> export.export(jax.jit(new_prim.bind))(1.)
Traceback (most recent call last):
ValueError: Cannot serialize code with custom calls whose targets have no compatibility guarantees: my_new_bind
>>> # We can avoid the error if we pass a `DisabledSafetyCheck.custom_call`
>>> exp = export.export(
... jax.jit(new_prim.bind),
... disabled_checks=[export.DisabledSafetyCheck.custom_call("my_new_prim")])(1.)
有关确保兼容性的开发者信息，请参阅 确保前向和后向兼容性。
跨平台和多平台导出#
JAX 降级对于少数 JAX 基元是平台特定的。默认情况下，代码是针对导出机器上存在的加速器进行降级和导出的。
>>> from jax import export
>>> export.default_export_platform()
'cpu'
存在一个安全检查，当尝试在没有导出代码时所用的加速器的机器上编译 Exported 对象时，会引发错误。
您可以明确指定应为哪些平台导出代码。这允许您指定与导出时可用的加速器不同的加速器，甚至可以指定多平台导出，以获得可在多个平台上编译和执行的 Exported 对象。
>>> import jax
>>> from jax import export
>>> from jax import lax
>>> # You can specify the export platform, e.g., `tpu`, `cpu`, `cuda`, `rocm`
>>> # even if the current machine does not have that accelerator.
>>> exp = export.export(jax.jit(lax.cos), platforms=['tpu'])(1.)
>>> # But you will get an error if you try to compile `exp`
>>> # on a machine that does not have TPUs.
>>> exp.call(1.)
Traceback (most recent call last):
ValueError: Function 'cos' was lowered for platforms '('tpu',)' but it is used on '('cpu',)'.
>>> # We can avoid the error if we pass a `DisabledSafetyCheck.platform`
>>> # parameter to `export`, e.g., because you have reasons to believe
>>> # that the code lowered will run adequately on the current
>>> # compilation platform (which is the case for `cos` in this
>>> # example):
>>> exp_unsafe = export.export(jax.jit(lax.cos),
... platforms=['tpu'],
... disabled_checks=[export.DisabledSafetyCheck.platform()])(1.)
>>> exp_unsafe.call(1.)
Array(0.5403023, dtype=float32, weak_type=True)
# and similarly with multi-platform lowering
>>> exp_multi = export.export(jax.jit(lax.cos),
... platforms=['tpu', 'cpu', 'cuda'])(1.)
>>> exp_multi.call(1.)
Array(0.5403023, dtype=float32, weak_type=True)
对于多平台导出，StableHLO 将包含多个降级版本，但仅针对那些需要它的基元，因此生成的模块大小应仅比默认导出的模块略大。在极端情况下，当序列化没有任何需要平台特定降级的基元的模块时，您将获得与单平台导出相同的 StableHLO。
>>> import jax
>>> from jax import export
>>> from jax import lax
>>> # A largish function
>>> def f(x):
... for i in range(1000):
... x = jnp.cos(x)
... return x
>>> exp_single = export.export(jax.jit(f))(1.)
>>> len(exp_single.mlir_module_serialized)
9220
>>> exp_multi = export.export(jax.jit(f),
... platforms=["cpu", "tpu", "cuda"])(1.)
>>> len(exp_multi.mlir_module_serialized)
9282
形状多态导出#
在 JIT 模式下使用时，JAX 会分别为每种输入形状组合追踪和降级函数。导出时，在某些情况下可以使用维度变量来处理某些输入维度，从而获得可用于多种输入形状组合的导出工件。
请参阅 形状多态性 (Shape polymorphism) 文档。
设备多态导出#
导出工件可能包含用于输入、输出和某些中间件的分片注释，但这些注释并不直接指向导出时存在的实际物理设备。相反，分片注释指向逻辑设备。这意味着您可以在与导出时不同的物理设备上编译和运行导出的工件。
实现设备多态导出的最清晰方法是使用由 jax.sharding.AbstractMesh 构建的分片，该分片仅包含网格形状和轴名称。但是，如果您使用为具有具体设备的网格构建的分片，也可以达到相同的效果，因为网格中的实际设备在追踪和降级时会被忽略。
>>> import jax
>>> from jax import export
>>> from jax.sharding import AbstractMesh, Mesh, NamedSharding
>>> from jax.sharding import PartitionSpec as P
>>>
>>> # Use an AbstractMesh for exporting
>>> export_mesh = AbstractMesh((4,), ("a",))
>>> def f(x):
... return x.T
>>> exp = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((32,), dtype=np.int32,
... sharding=NamedSharding(export_mesh, P("a"))))
>>> # `exp` knows for how many devices it was exported.
>>> exp.nr_devices
4
>>> # and it knows the shardings for the inputs. These will be applied
>>> # when the exported is called.
>>> exp.in_shardings_jax(export_mesh)
(NamedSharding(mesh=AbstractMesh('a': 4, axis_types=(Auto,)), spec=P('a',)),)
>>> # You can also use a concrete set of devices for exporting
>>> concrete_devices = jax.local_devices()[:4]
>>> concrete_mesh = Mesh(concrete_devices, ("a",))
>>> exp2 = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((32,), dtype=np.int32,
... sharding=NamedSharding(concrete_mesh, P("a"))))
>>> # When you call an Exported, you must use a concrete set of devices
>>> arg = jax.device_put(jnp.arange(8 * 4),
... NamedSharding(concrete_mesh, P("a")))
>>> res1 = exp.call(arg)
>>> # Check out the first 2 shards of the result
>>> [f"device={s.device} index={s.index}" for s in res1.addressable_shards[:2]]
['device=cpu:0 index=(slice(0, 8, None),)', 'device=cpu:1 index=(slice(8, 16, None),)']
>>> # We can call `exp` with some other 4 devices and another
>>> # mesh with a different shape, as long as the number of devices is
>>> # the same.
>>> other_mesh = Mesh(np.array(jax.local_devices()[2:6]).reshape((2, 2)), ("b", "c"))
>>> res2 = exp.call(jax.device_put(arg,
... NamedSharding(other_mesh, P("b"))))
>>> # Check out the first 2 shards of the result. Notice that the output is
>>> # sharded similarly; this means that the input was resharded according to the
>>> # exp.in_shardings.
>>> [f"device={s.device} index={s.index}" for s in res2.addressable_shards[:2]]
['device=cpu:2 index=(slice(0, 8, None),)', 'device=cpu:3 index=(slice(8, 16, None),)']
尝试使用与导出时不同数量的设备来调用导出工件是错误的。
>>> import jax
>>> from jax import export
>>> from jax.sharding import Mesh, NamedSharding
>>> from jax.sharding import PartitionSpec as P
>>> export_devices = jax.local_devices()
>>> export_mesh = Mesh(np.array(export_devices), ("a",))
>>> def f(x):
... return x.T
>>> exp = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((4 * len(export_devices),), dtype=np.int32,
... sharding=NamedSharding(export_mesh, P("a"))))
>>> arg = jnp.arange(4 * len(export_devices))
>>> exp.call(arg)
Traceback (most recent call last):
ValueError: Exported module f was lowered for 8 devices and is called in a context with 1 devices. This is disallowed because: the module was lowered for more than 1 device.
有一些辅助函数可以使用在调用点构建的新网格来对导出工件的输入进行分片。
>>> import jax
>>> from jax import export
>>> from jax.sharding import Mesh, NamedSharding
>>> from jax.sharding import PartitionSpec as P
>>> export_devices = jax.local_devices()
>>> export_mesh = Mesh(np.array(export_devices), ("a",))
>>> def f(x):
... return x.T
>>> exp = export.export(jax.jit(f))(
... jax.ShapeDtypeStruct((4 * len(export_devices),), dtype=np.int32,
... sharding=NamedSharding(export_mesh, P("a"))))
>>> # Prepare the mesh for calling `exp`.
>>> calling_mesh = Mesh(np.array(export_devices[::-1]), ("a",))
>>> # Shard the arg according to what `exp` expects.
>>> arg = jnp.arange(4 * len(export_devices))
>>> sharded_arg = jax.device_put(arg, exp.in_shardings_jax(calling_mesh)[0])
>>> res = exp.call(sharded_arg)
作为一种特殊功能，如果函数是为 1 个设备导出的，并且不包含任何分片注释，则可以在形状相同但分片在多个设备上的参数上调用它，编译器将适当地对该函数进行分片。
>>> import jax
>>> from jax import export
>>> from jax.sharding import Mesh, NamedSharding
>>> from jax.sharding import PartitionSpec as P
>>> def f(x):
... return jnp.cos(x)
>>> arg = jnp.arange(4)
>>> exp = export.export(jax.jit(f))(arg)
>>> exp.in_avals
(ShapedArray(int32[4]),)
>>> exp.nr_devices
1
>>> # Prepare the mesh for calling `exp`.
>>> calling_mesh = Mesh(jax.local_devices()[:4], ("b",))
>>> # Shard the arg according to what `exp` expects.
>>> sharded_arg = jax.device_put(arg,
... NamedSharding(calling_mesh, P("b")))
>>> res = exp.call(sharded_arg)
调用约定版本#
JAX 导出支持随着时间的推移而演变，例如为了支持副作用 (effects)。为了支持兼容性（请参阅 兼容性保证），我们为每个 Exported 维护一个调用约定版本。截至 2024 年 6 月，所有函数均以第 9 版导出（最新版本，请参阅 所有调用约定版本）。
>>> from jax import export
>>> exp: export.Exported = export.export(jnp.cos)(1.)
>>> exp.calling_convention_version
10
在任何特定时间，导出 API 可能支持一系列调用约定版本。您可以使用 --jax_export_calling_convention_version 标志或 JAX_EXPORT_CALLING_CONVENTION_VERSION 环境变量来控制使用哪个调用约定版本。
>>> from jax import export
>>> (export.minimum_supported_calling_convention_version, export.maximum_supported_calling_convention_version)
(9, 10)
>>> from jax._src import config
>>> with config.jax_export_calling_convention_version(10):
... exp = export.export(jnp.cos)(1.)
... exp.calling_convention_version
10
我们保留删除对生成或消耗超过 6 个月旧的调用约定版本的支持的权利。
内部函数使用不同的调用约定：可选的平台索引参数、可选的维度变量参数（i32 或 i64 类型的标量张量），随后是可选的 token 参数（在存在有序副作用时），然后是常规数组参数。维度参数对应于 考虑将一个具有一个 实际计算位于 在调用约定第 9 版之前，副作用的调用约定有所不同： 同样从调用约定第 9 版开始，包含平台索引或维度变量值的函数参数具有一个 请注意， 此处列出了调用约定版本的历史记录： 第 1 版使用 MHLO 和 CHLO 来序列化代码，不再支持。 第 2 版支持 StableHLO 和 CHLO。自 2022 年 10 月起使用。不再支持。 第 3 版支持平台检查和多个平台。自 2023 年 2 月起使用。不再支持。 第 4 版支持具有兼容性保证的 StableHLO。这是 JAX 本地序列化启动时最早的版本。自 2023 年 3 月 15 日起在 JAX 中使用 (cl/516885716)。从 2023 年 3 月 28 日起，我们停止使用 第 5 版增加了对 第 6 版增加了对 第 7 版增加了对 第 8 版增加了对 第 9 版增加了对副作用的支持。有关精确的调用约定，请参阅 第 10 版将 模块调用约定#
Exported.mlir_module 具有一个 main 函数，如果模块支持多个平台（len(platforms) > 1），它接受一个可选的第一个平台索引参数，随后是对应于有序副作用的 token 参数，然后是保留的数组参数（对应于 module_kept_var_idx 和 in_avals）。平台索引是一个 i32 或 i64 标量，将当前编译平台的索引编码为 platforms 序列。args_avals 中出现的维度变量，按其名称排序。f32[w, 2 * h] 类型数组参数的函数进行降级，其中 w 和 h 是两个维度变量。假设我们使用多平台降级，并且有一个有序副作用。main 函数如下所示： func public main(
platform_index: i32 {jax.global_constant="_platform_index"},
token_in: token,
arg: f32[?, ?]) {
arg_w = hlo.get_dimension_size(arg, 0)
dim1 = hlo.get_dimension_size(arg, 1)
arg_h = hlo.floordiv(dim1, 2)
call _check_shape_assertions(arg) # See below
token = new_token()
token_out, res = call _wrapped_jax_export_main(platform_index,
arg_h,
arg_w,
token_in,
arg)
return token_out, res
}
_wrapped_jax_export_main 中，同时接受 h 和 w 维度变量的值。_wrapped_jax_export_main 的签名为： func private _wrapped_jax_export_main(
platform_index: i32 {jax.global_constant="_platform_index"},
arg_h: i32 {jax.global_constant="h"},
arg_w: i32 {jax.global_constant="w"},
arg_token: stablehlo.token {jax.token=True},
arg: f32[?, ?]) -> (stablehlo.token, ...)
main 函数不接受或返回 token。相反，该函数创建 i1[0] 类型的虚拟 token 并将其传递给 _wrapped_jax_export_main。_wrapped_jax_export_main 接受 i1[0] 类型的虚拟 token，并将在内部创建真实的 token 以传递给内部函数。内部函数使用真实的 token（在调用约定第 9 版之前和之后都是如此）。jax.global_constant 字符串属性，其值为全局常量的名称，即 _platform_index 或维度变量名。如果名称未知，则全局常量名可能为空。某些全局常量计算使用内部函数，例如 floor_divide。此类函数的参数对于所有属性都有一个 jax.global_constant 属性，这意味着函数的结果也是一个全局常量。main 包含对 _check_shape_assertions 的调用。JAX 追踪假设 arg.shape[1] 是偶数，且 w 和 h 的值均 >= 1。我们在调用模块时必须检查这些约束。我们使用一个特殊的自定义调用 @shape_assertion，它接受一个布尔值作为第一个操作数，一个可能包含格式说明符 {0}, {1}, … 的字符串 error_message 属性，以及与格式说明符相对应的可变数量的整数标量操作数。 func private _check_shape_assertions(arg: f32[?, ?]) {
# Check that w is >= 1
arg_w = hlo.get_dimension_size(arg, 0)
custom_call @shape_assertion(arg_w >= 1, arg_w,
error_message="Dimension variable 'w' must have integer value >= 1. Found {0}")
# Check that dim1 is even
dim1 = hlo.get_dimension_size(arg, 1)
custom_call @shape_assertion(dim1 % 2 == 0, dim1 % 2,
error_message="Division had remainder {0} when computing the value of 'h')
# Check that h >= 1
arg_h = hlo.floordiv(dim1, 2)
custom_call @shape_assertion(arg_h >= 1, arg_h,
error_message=""Dimension variable 'h' must have integer value >= 1. Found {0}")
调用约定版本#
dim_args_spec (cl/520033493)。该版本的支持于 2023 年 10 月 17 日结束 (cl/573858283)。call_tf_graph 的支持。目前用于某些特定用例。自 2023 年 5 月 3 日起在 JAX 中使用 (cl/529106145)。disabled_checks 属性的支持。此版本强制要求 platforms 属性不能为空。自 2023 年 6 月 7 日起由 XlaCallModule 支持，并自 2023 年 6 月 13 日起在 JAX (JAX 0.4.13) 中可用。stablehlo.shape_assertion 操作以及 disabled_checks 中指定的 shape_assertions 的支持。请参阅 形状多态性下的错误。自 2023 年 7 月 12 日起由 XlaCallModule 支持，自 2023 年 7 月 20 日起在 JAX 序列化中可用 (JAX 0.4.14)，并自 2023 年 8 月 12 日起成为默认设置 (JAX 0.4.15)。jax.uses_shape_polymorphism 模块属性的支持，并仅在存在该属性时启用形状优化通道 (shape refinement pass)。自 2023 年 7 月 21 日起由 XlaCallModule 支持，自 2023 年 7 月 26 日起在 JAX 中可用 (JAX 0.4.14)，并自 2023 年 10 月 21 日起成为默认设置 (JAX 0.4.20)。export.Exported 的文档字符串。在此调用约定版本中，我们还用 jax.global_constant 属性标记了平台索引和维度变量参数。自 2023 年 10 月 27 日起由 XlaCallModule 支持，自 2023 年 10 月 20 日起在 JAX 中可用 (JAX 0.4.20)，并自 2024 年 2 月 1 日起成为默认设置 (JAX 0.4.24)。截至 2024 年 3 月 27 日，这是唯一受支持的版本。jax.config.use_shardy_partitioner 的值传播到 XlaCallModule。自 2025 年 5 月 20 日起由 XlaCallModule 支持，并自 2025 年 7 月 14 日起在 JAX 中成为默认设置 (JAX 0.7.0)。
开发者文档#
调试#
您可以记录导出的模块，在开源 (OSS) 和 Google 内部有稍微不同的标志。在 OSS 中，您可以执行以下操作：
# Log from python
python tests/export_test.py JaxExportTest.test_basic -v=3
# Or, log from pytest to /tmp/mylog.txt
pytest tests/export_test.py -k test_basic --log-level=3 --log-file=/tmp/mylog.txt
您将看到格式如下的日志行：
I0619 10:54:18.978733 8299482112 _export.py:606] Exported JAX function: fun_name=sin version=9 lowering_platforms=('cpu',) disabled_checks=()
I0619 10:54:18.978767 8299482112 _export.py:607] Define JAX_DUMP_IR_TO to dump the module.
如果您将环境变量 JAX_DUMP_IR_TO 设置为某个目录，导出的（以及 JIT 编译的）HLO 模块将保存在该目录中。
JAX_DUMP_IR_TO=/tmp/export.dumps pytest tests/export_test.py -k test_basic --log-level=3 --log-file=/tmp/mylog.txt
INFO absl:_export.py:606 Exported JAX function: fun_name=sin version=9 lowering_platforms=('cpu',) disabled_checks=()
INFO absl:_export.py:607 The module was dumped to jax_ir0_jit_sin_export.mlir.
您将看到导出的模块（命名为 ..._export.mlir）和 JIT 编译的模块（命名为 ..._compile.mlir）。
$ ls -l /tmp/export.dumps/
total 32
-rw-rw-r--@ 1 necula wheel 2316 Jun 19 11:04 jax_ir0_jit_sin_export.mlir
-rw-rw-r--@ 1 necula wheel 2279 Jun 19 11:04 jax_ir1_jit_sin_compile.mlir
-rw-rw-r--@ 1 necula wheel 3377 Jun 19 11:04 jax_ir2_jit_call_exported_compile.mlir
-rw-rw-r--@ 1 necula wheel 2333 Jun 19 11:04 jax_ir3_jit_my_fun_export.mlir
设置 JAX_DEBUG_LOG_MODULES=jax._src.export 以启用额外的调试日志记录。
确保前向和后向兼容性#
本节讨论 JAX 开发者应采取的流程，以确保兼容性保证。
一个复杂之处在于，外部用户将 JAX 和 jaxlib 安装在单独的包中，并且用户通常最终使用比 JAX 更旧的 jaxlib。我们观察到自定义调用存在于 jaxlib 中，而对于导出工件的消费者而言，只有 jaxlib 是相关的。为了简化流程，我们向外部用户明确期望：兼容性窗口是根据 jaxlib 版本定义的，即使 JAX 可能使用旧版本运行，用户也有责任确保他们使用新的 jaxlib 进行导出。
因此，我们只关心 jaxlib 版本。我们可以在发布 jaxlib 版本时启动后向兼容性弃用时钟，即使我们不强制要求它是允许的最低版本。
假设我们需要添加、删除或更改 JAX 降级规则使用的自定义调用目标 T 的语义。以下是可能的编年史（针对驻留在 jaxlib 中的自定义调用目标的更改）：
第 “D - 1” 天，变更前。假设当前激活的内部 JAX 版本为
0.4.31（下一个 JAX 和 jaxlib 版本的版本号）。JAX 降级规则使用自定义调用T。第 “D” 天，我们添加新的自定义调用目标
T_NEW。我们应该创建一个新的自定义调用目标，并在大约 6 个月后清理旧目标，而不是原地更新T。请参阅示例 PR #20997，其中实现了以下步骤。
我们添加自定义调用目标
T_NEW。我们更改之前使用
T的 JAX 降级规则，使其有条件地使用T_NEW，如下所示：
from jax._src import config from jax._src.lib import version as jaxlib_version def my_lowering_rule(ctx: LoweringRuleContext, ...): if ctx.is_forward_compat() or jaxlib_version < (0, 4, 31): # this is the old lowering, using target T, while we # are in forward compatibility mode for T, or we # are in OSS and are using an old jaxlib. return hlo.custom_call("T", ...) else: # This is the new lowering, using target T_NEW, for # when we use a jaxlib with version `>= (0, 4, 31)` # (or when this is internal usage), and also we are # in JIT mode. return hlo.custom_call("T_NEW", ...)
请注意，在 JIT 模式下，或者如果用户传递
--jax_export_ignore_forward_compatibility=true，前向兼容模式始终为 false。请注意，此时导出仍然不会使用
T_NEW。
这可以在上一步之后的任何时间完成，且必须在下一步之前完成：为
T_NEW添加后向兼容性测试，并将T_NEW添加到_export.py中的_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE列表中。关于添加后向兼容性测试的说明位于 export_back_compat_test_util.py 的顶部。
示例见 PR #29488。
请注意，如果您在下一步之前执行此操作，导出仍然不会使用
T_NEW降级，并且您必须在self.run_one_test的调用周围添加with config.export_ignore_forward_compatibility(True):。当您真正进入第 4 步时，可以将其删除。您可能还需要仅针对新版本的 jaxlib 启用测试。
第 “D + 21” 天（前向兼容性窗口结束；甚至可以晚于 21 天）：我们删除降级代码中的
forward_compat_mode，因此现在只要我们使用新的jaxlib，导出就会开始使用新的自定义调用目标T_NEW。第 “RELEASE > D” 天（D 之后第一个 JAX 发布日期，当我们发布版本
0.4.31时）：我们开始计算 6 个月的后向兼容性时钟。请注意，仅当T属于我们已经保证稳定性的自定义调用目标时，即列在_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE中时，这才是相关的。如果
RELEASE处于前向兼容性窗口[D, D + 21]中，并且如果我们使RELEASE成为允许的最低 jaxlib 版本，那么我们可以删除 JIT 分支中的jaxlib_version < (0, 4, 31)条件。
第 “RELEASE + 180” 天（后向兼容性窗口结束，甚至可以晚于 180 天）：此时，我们必须已经提高了最低 jaxlib 要求，以便降级条件
jaxlib_version < (0, 4, 31)已被删除，并且 JAX 降级无法再生成对T的自定义调用。我们删除旧自定义调用目标
T的 C++ 实现。我们也删除
T的后向兼容性测试。
从 jax.experimental.export 迁移指南#
2024 年 6 月 18 日（JAX 版本 0.4.30），我们弃用了 jax.experimental.export API，转而使用 jax.export API。有一些细微的变化：
jax.experimental.export.export:旧函数曾经允许传入任何 Python 可调用对象或
jax.jit的结果。现在只接受后者。在调用export之前，您必须手动对要导出的函数应用jax.jit。旧的
lowering_parameters关键字参数现在命名为platforms。
jax.experimental.export.default_lowering_platform()现在位于jax.export.default_export_platform()。jax.experimental.export.call现在是jax.export.Exported对象的一个方法。您应该使用exp.call而不是export.call(exp)。jax.experimental.export.serialize现在是jax.export.Exported对象的一个方法。您应该使用exp.serialize()而不是export.serialize(exp)。配置标志
--jax-serialization-version已弃用。请使用--jax-export-calling-convention-version。jax.experimental.export.minimum_supported_serialization_version的值现在位于jax.export.minimum_supported_calling_convention_version。jax.export.Exported的以下字段已重命名：uses_shape_polymorphism现在是uses_global_constants。mlir_module_serialization_version现在是calling_convention_version。lowering_platforms现在是platforms。