导出和序列化阶段输出计算#
The 提前降低和编译 APIs 生成的对象可用于调试或在同一进程中进行编译和执行。有时，您可能希望将一个降低的 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 的互操作。
支持反向模式自动微分#
序列化可以选择性地支持高阶反向模式自动微分。这是通过序列化原始函数的 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()。
请注意，重新水合的函数不支持任何其他变换，例如，前向模式自动微分（jvp），或 jax.vmap()。
兼容性保证#
出于几个原因，你不应该将仅通过降低 (jax.jit(f).lower(1.).compiler_ir()) 获得的原始 StableHLO 用于存档或在另一个进程中进行编译。
首先，编译可能会使用不同版本的编译器，支持不同版本的 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 支持（cl/547482522），自 2023 年 7 月 20 日起在 JAX 序列化中可用（JAX 0.4.14），自 2023 年 8 月 12 日起为默认设置（JAX 0.4.15）。版本 8 增加了对
jax.uses_shape_polymorphism模块属性的支持，并且仅在属性存在时启用形状细化传递。自 2023 年 7 月 21 日（cl/549973693）起由 XlaCallModule 支持，自 2023 年 7 月 26 日（JAX 0.4.14）起在 JAX 中可用，自 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）和谷歌中使用略有不同的标志来记录导出的模块。在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 设置为一个目录，导出的（以及即时编译的）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）和即时编译的模块（命名为 ..._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, ...): lowering_parameters = ctx.module_context.lowering_parameters forward_compat_mode = (lowering_parameters.for_export and not lowering_parameters.export_ignore_forward_compatibility) if forward_compat_mode 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添加了一个向后兼容性测试。
Day “RELEASE > D” (在
D之后的第一个 JAX 发布日期，我们发布版本0.4.31): 我们开始为期 6 个月的向后兼容性时钟。请注意，这仅在与我们已经保证稳定性的自定义调用目标相关时才适用，即，这些目标列在_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_constantsmlir_module_serialization_version现在是calling_convention_versionlowering_platforms现在为platforms。