导出和序列化分阶段计算

内容

导出和序列化分阶段计算#

The 提前降低和编译 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
>>> exported.fun_name
'f'

>>> exported.in_avals
(ShapedArray(float32[]),)

>>> print(re.search(r".*@main.*", exported.mlir_module()).group(0))
  func.func public @main(%arg0: tensor<f32> {mhlo.layout_mode = "default"} loc("x")) -> (tensor<f32> {jax.result_info = "", mhlo.layout_mode = "default"}) {

>>> # 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)
>>> rehydrated_exp.in_avals
(ShapedArray(float32[]),)

>>> def callee(y):
...  return 3. * rehydrated_exp.call(y * 4.)

>>> callee(1.)
Array(96., dtype=float32)

序列化分为两个阶段

导出以生成一个包含已降低函数的 StableHLO 以及从另一个 JAX 函数调用它所需元数据的 jax.export.Exported 对象。我们计划添加代码以从 TensorFlow 生成 Exported 对象，以及使用来自 TensorFlow 和 PyTorch 的 Exported 对象。
使用 flatbuffers 格式将数据实际序列化为字节数组。有关可用于与 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()。

兼容性保证#

出于多种原因，您不应将仅通过降低获得的原始 StableHLO（jax.jit(f).lower(1.).compiler_ir()）用于归档和在另一个进程中进行编译。

首先，编译可能使用不同版本的编译器，支持不同版本的 StableHLO。 jax.export 模块通过使用 StableHLO 的可移植工件特性来处理 StableHLO 操作集的可能演变，从而解决了此问题。

自定义调用的兼容性保证#

其次，原始 StableHLO 可能包含引用 C++ 函数的自定义调用。JAX 使用自定义调用来降低少量原语，例如线性代数原语、分片注释或 Pallas 内核。这些不属于 StableHLO 的兼容性保证。这些函数的 C++ 实现很少更改，但可能会更改。

jax.export 提供以下导出兼容性保证：JAX 导出的工件可以由编译器和 JAX 运行时系统编译和执行，这些系统

**最多比用于导出 JAX 的版本新 6 个月**（我们说 JAX 导出提供**6 个月的向后兼容性**）。如果我们想将导出的工件存档以供以后编译和执行，这将很有用。
**最多比用于导出 JAX 的版本旧 3 周**（我们说 JAX 导出提供**3 周的前向兼容性**）。如果我们想使用在导出之前构建和部署的使用者（例如，在导出完成后已部署的推理系统）来编译和运行导出的工件，这将很有用。

（特定的兼容性窗口长度与 JAX 对 jax2tf 的承诺相同，并且基于 TensorFlow 兼容性。术语“向后兼容性”是从使用者的角度来看的，例如，推理系统。）

**导出和使用组件何时构建**很重要，而不是导出和编译何时发生。对于外部 JAX 用户，可以以不同的版本运行 JAX 和 jaxlib；重要的是 jaxlib 版本何时构建。

为了减少不兼容的可能性，内部 JAX 用户应

尽可能频繁地重建和重新部署使用者系统.

而外部用户则应

尽可能使用相同版本的 jaxlib 运行导出和使用者系统，以及
**使用最新发布的 jaxlib 版本**导出以供存档。

如果您绕过 jax.export API 获取 StableHLO 代码，则兼容性保证不适用。

为了确保前向兼容性，当我们更改 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> {mhlo.layout_mode = "default"}) -> (tensor<f32> {jax.result_info = "", mhlo.layout_mode = "default"}) {
    %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),
...    lowering_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),
...    lowering_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),
...                           lowering_platforms=["cpu", "tpu", "cuda"])(1.)
>>> len(exp_multi.mlir_module_serialized)  
9282

形状多态导出#

在 JIT 模式下使用时，JAX 将分别针对每种输入形状组合跟踪和降低函数。在导出时，在某些情况下可以使用维度变量来表示某些输入维度，以便获得可用于多种输入形状组合的导出工件。

请参阅形状多态性文档。

设备多态导出#

导出的工件可能包含输入、输出和某些中间体的分片注释，但这些注释不直接引用导出时存在的实际物理设备。相反，分片注释引用逻辑设备。这意味着您可以在用于导出的不同物理设备上编译和运行导出的工件。

>>> import jax
>>> from jax import export
>>> from jax.sharding import Mesh, NamedSharding
>>> from jax.sharding import PartitionSpec as P

>>> # Use the first 4 devices for exporting.
>>> export_devices = jax.local_devices()[:4]
>>> export_mesh = Mesh(export_devices, ("a",))
>>> def f(x):
...   return x.T

>>> arg = jnp.arange(8 * len(export_devices))
>>> exp = export.export(jax.jit(f, in_shardings=(NamedSharding(export_mesh, P("a")),)))(arg)

>>> # `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_hlo
({devices=[4]<=[4]},)

>>> res1 = exp.call(jax.device_put(arg,
...                                NamedSharding(export_mesh, P("a"))))

>>> # Check out the first 2 shards of the result
>>> [f"device={s.device} index={s.index}" for s in res1.addressable_shards[:2]]
['device=TFRT_CPU_0 index=(slice(0, 8, None),)',
 'device=TFRT_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=TFRT_CPU_2 index=(slice(0, 8, None),)',
 'device=TFRT_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

>>> arg = jnp.arange(4 * len(export_devices))
>>> exp = export.export(jax.jit(f, in_shardings=(NamedSharding(export_mesh, P("a")),)))(arg)

>>> 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

>>> arg = jnp.arange(4 * len(export_devices))
>>> exp = export.export(jax.jit(f, in_shardings=(NamedSharding(export_mesh, P("a")),)))(arg)

>>> # Prepare the mesh for calling `exp`.
>>> calling_mesh = Mesh(np.array(export_devices[::-1]), ("b",))

>>> # Shard the arg according to what `exp` expects.
>>> sharded_arg = jax.device_put(arg, exp.in_shardings_jax(calling_mesh)[0])
>>> res = exp.call(sharded_arg)

作为一个特殊的功能，如果一个函数是为 1 个设备导出的，并且它不包含任何分片注释，那么它可以在具有相同形状但分布在多个设备上的参数上调用，并且编译器将适当地分片函数

```python
>>> 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 导出支持随着时间的推移而发展，例如，以支持效果。为了支持兼容性（请参阅兼容性保证），我们为每个 Exported 保留调用约定版本。截至 2024 年 6 月，所有使用版本 9（最新版本，请参阅所有调用约定版本）导出的函数

>>> from jax import export
>>> exp: export.Exported = export.export(jnp.cos)(1.)
>>> exp.calling_convention_version
9

在任何给定时间，导出 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, 9)

>>> from jax._src import config
>>> with config.jax_export_calling_convention_version(9):
...  exp = export.export(jnp.cos)(1.)
...  exp.calling_convention_version
9

我们保留删除对生成或使用 6 个月以上旧的调用约定版本的支持的权利。

模块调用约定#

Exported.mlir_module 具有一个 main 函数，如果模块支持多个平台（len(platforms) > 1），则该函数将采用可选的第一个平台索引参数，后跟对应于有序效果的令牌参数，然后是保留的数组参数（对应于 module_kept_var_idx 和 in_avals）。平台索引是 i32 或 i64 标量，用于将当前编译平台的索引编码到 platforms 序列中。

内部函数使用不同的调用约定：一个可选的平台索引参数，可选的维度变量参数（类型为 i32 或 i64 的标量张量），然后是可选的令牌参数（在存在有序效应的情况下），最后是常规的数组参数。维度参数对应于出现在 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, ...)

在调用约定版本 9 之前，效应的调用约定是不同的： main 函数不接收或返回令牌。相反，该函数创建类型为 i1[0] 的虚拟令牌并将其传递给 _wrapped_jax_export_main。 _wrapped_jax_export_main 接收类型为 i1[0] 的虚拟令牌，并在内部创建真实的令牌传递给内部函数。内部函数使用真实的令牌（在调用约定版本 9 之前和之后）

同样，从调用约定版本 9 开始，包含平台索引或维度变量值的函数参数具有一个 jax.global_constant 字符串属性，其值为全局常量的名称，即 _platform_index 或维度变量名称。如果未知，则全局常量名称可以为空。一些全局常量计算使用内部函数，例如，对于 floor_divide。此类函数的参数对所有属性都有一个 jax.global_constant 属性，这意味着函数的结果也是一个全局常量。

请注意， main 包含对 _check_shape_assertions 的调用。JAX 跟踪假设 arg.shape[1] 是偶数，并且 w 和 h 的值都 >= 1。当我们调用模块时，必须检查这些约束。我们使用一个特殊的自定义调用 @shape_assertion，它接收一个布尔类型的第一个操作数，一个字符串 error_message 属性，该属性可能包含格式说明符 {0}、 {1} 等，以及与格式说明符相对应的可变数量的整数标量操作数。

       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}")

调用约定版本#

我们在此列出调用约定版本号的历史记录

版本 1 使用 MHLO & CHLO 序列化代码，不再支持。
版本 2 支持 StableHLO & CHLO。从 2022 年 10 月开始使用。不再支持。
版本 3 支持平台检查和多个平台。从 2023 年 2 月开始使用。不再支持。
版本 4 支持具有兼容性保证的 StableHLO。这是 JAX 原生序列化启动时的最早版本。从 2023 年 3 月 15 日开始在 JAX 中使用（cl/516885716）。从 2023 年 3 月 28 日开始，我们停止使用 dim_args_spec（cl/520033493）。对该版本的支持于 2023 年 10 月 17 日停止（cl/573858283）。
版本 5 添加了对 call_tf_graph 的支持。这目前用于某些特殊用例。从 2023 年 5 月 3 日开始在 JAX 中使用（cl/529106145）。
版本 6 添加了对 disabled_checks 属性的支持。此版本要求 platforms 属性非空。从 2023 年 6 月 7 日起由 XlaCallModule 支持，并从 2023 年 6 月 13 日起在 JAX 中可用（JAX 0.4.13）。
版本 7 添加了对 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）。
版本 8 添加了对 jax.uses_shape_polymorphism 模块属性的支持，并且仅在存在该属性时启用形状细化传递。从 2023 年 7 月 21 日起由 XlaCallModule 支持，从 2023 年 7 月 26 日起在 JAX 中可用（JAX 0.4.14），并从 2023 年 10 月 21 日起成为默认值（JAX 0.4.20）。
版本 9 添加了对效应的支持。有关精确的调用约定，请参阅 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 日，这是唯一受支持的版本。

开发者文档#

调试#

您可以使用不同的标志在 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

在 Google 内部，您可以通过使用 --vmodule 参数指定不同模块的日志级别来打开日志记录，例如 --vmodule=_export=3。

确保向前和向后兼容性#

本节讨论 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 时，前向兼容模式始终为假
我们将 T_NEW 添加到 _CUSTOM_CALL_TARGETS_GUARANTEED_STABLE 的列表中，位于 _export.py 中。

“D + 21”日（前向兼容性窗口结束；可能晚于 21 天）：我们删除了降低代码中的 forward_compat_mode，因此现在只要我们使用新的 jaxlib，导出就会开始使用新的自定义调用目标 T_NEW。
- 我们为 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 版本提升，以便 JIT 分支中的 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 的结果。现在仅接受后者。您必须手动将 jax.jit 应用于要导出的函数，然后再调用 export。
- 旧的 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。