JAX 中文文档（四）（4）https://developer.aliyun.com/article/1559796
兼容性保证
您不应仅从降低中获取的原始 StableHLO（jax.jit(f).lower(1.).compiler_ir()）用于归档和在另一个进程中进行编译，有几个原因。
首先，编译可能使用不同版本的编译器，支持不同版本的 StableHLO。jax.export 模块通过使用 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 代码，则不适用兼容性保证。
只有部分自定义调用被保证稳定，并具有兼容性保证（参见列表）。我们会持续向允许列表中添加更多自定义调用目标，同时进行向后兼容性测试。如果尝试序列化调用其他自定义调用目标的代码，则在导出期间会收到错误。
如果您希望禁用特定自定义调用的此安全检查，例如目标为 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} : (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 >>> # 准备用于调用 `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) ```py ## Calling convention versions The JAX export support has evolved over time, e.g., to support effects. In order to support compatibility (see compatibility guarantees) we maintain a calling convention version for each `Exported`. As of June 2024, all function exported with version 9 (the latest, see all calling convention versions):
from jax import export
exp: export.Exported = export.export(jnp.cos)(1.)
exp.calling_convention_version
9
At any given time, the export APIs may support a range of calling convention versions. You can control which calling convention version to use using the `--jax-export-calling-convention-version` flag or the `JAX_EXPORT_CALLING_CONVENTION_VERSION` environment variable:
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
We reserve the right to remove support for generating or consuming calling convention versions older than 6 months. ### Module calling convention The `Exported.mlir_module` has a `main` function that takes an optional first platform index argument if the module supports multiple platforms (`len(platforms) > 1`), followed by the token arguments corresponding to the ordered effects, followed by the kept array arguments (corresponding to `module_kept_var_idx` and `in_avals`). The platform index is a i32 or i64 scalar encoding the index of the current compilation platform into the `platforms` sequence. Inner functions use a different calling convention: an optional platform index argument, optional dimension variable arguments (scalar tensors of type i32 or i64), followed by optional token arguments (in presence of ordered effects), followed by the regular array arguments. The dimension arguments correspond to the dimension variables appearing in the `args_avals`, in sorted order of their names. Consider the lowering of a function with one array argument of type `f32[w, 2 * h]`, where `w` and `h` are two dimension variables. Assume that we use multi-platform lowering, and we have one ordered effect. The `main` function will be as follows:
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 }
The actual computation is in `_wrapped_jax_export_main`, taking also the values of `h` and `w` dimension variables. The signature of the `_wrapped_jax_export_main` is:
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, ...)
Prior to calling convention version 9 the calling convention for effects was different: the `main` function does not take or return a token. Instead the function creates dummy tokens of type `i1[0]` and passes them to the `_wrapped_jax_export_main`. The `_wrapped_jax_export_main` takes dummy tokens of type `i1[0]` and will create internally real tokens to pass to the inner functions. The inner functions use real tokens (both before and after calling convention version 9) Also starting with calling convention version 9, function arguments that contain the platform index or the dimension variable values have a `jax.global_constant` string attribute whose value is the name of the global constant, either `_platform_index` or a dimension variable name. The global constant name may be empty if it is not known. Some global constant computations use inner functions, e.g., for `floor_divide`. The arguments of such functions have a `jax.global_constant` attribute for all attributes, meaning that the result of the function is also a global constant. Note that `main` contains a call to `_check_shape_assertions`. JAX tracing assumes that `arg.shape[1]` is even, and that both `w` and `h` have values >= 1\. We must check these constraints when we invoke the module. We use a special custom call `@shape_assertion` that takes a boolean first operand, a string `error_message` attribute that may contain format specifiers `{0}`, `{1}`, …, and a variadic number of integer scalar operands corresponding to the format specifiers.
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, error_message="Dimension variable 'h' must have integer value >= 1\. Found non-zero remainder {0}") # 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}")
### Calling convention versions We list here a history of the calling convention version numbers: + Version 1 used MHLO & CHLO to serialize the code, not supported anymore. + Version 2 supports StableHLO & CHLO. Used from October 2022\. Not supported anymore. + Version 3 supports platform checking and multiple platforms. Used from February 2023\. Not supported anymore. + Version 4 supports StableHLO with compatibility guarantees. This is the earliest version at the time of the JAX native serialization launch. Used in JAX from March 15, 2023 (cl/516885716). Starting with March 28th, 2023 we stopped using `dim_args_spec` (cl/520033493). The support for this version was dropped on October 17th, 2023 (cl/573858283). + Version 5 adds support for `call_tf_graph`. This is currently used for some specialized use cases. Used in JAX from May 3rd, 2023 (cl/529106145). + 第 6 版添加了对 `disabled_checks` 属性的支持。此版本要求 `platforms` 属性不为空。自 2023 年 6 月 7 日由 XlaCallModule 支持，自 2023 年 6 月 13 日（JAX 0.4.13）起支持 JAX。 + 第 7 版增加了对 `stablehlo.shape_assertion` 操作和在 `disabled_checks` 中指定的 `shape_assertions` 的支持。参见[形状多态性存在错误](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#errors-in-presence-of-shape-polymorphism)。自 2023 年 7 月 12 日（cl/547482522）由 XlaCallModule 支持，自 2023 年 7 月 20 日（JAX 0.4.14）起支持 JAX 序列化，并自 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 版添加了对 effects 的支持。详见 `export.Exported` 的文档字符串获取准确的调用约定。在此调用约定版本中，我们还使用 `jax.global_constant` 属性标记平台索引和维度变量参数。自 2023 年 10 月 27 日由 XlaCallModule 支持，自 2023 年 10 月 20 日（JAX 0.4.20）起支持 JAX，并自 2024 年 2 月 1 日（JAX 0.4.24）起成为默认选项。截至 2024 年 3 月 27 日，这是唯一支持的版本。 ## 从 `jax.experimental.export` 迁移指南。 在 2024 年 6 月 14 日，我们废弃了 `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` 对象的一个方法。不再使用 `export.call(exp)`，应使用 `exp.call`。 + `jax.experimental.export.serialize` 现在是 `jax.export.Exported` 对象的一个方法。不再使用 `export.serialize(exp)`，应使用 `exp.serialize()`。 + 配置标志 `--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`。 2023 (cl/529106145). + 第 6 版添加了对 `disabled_checks` 属性的支持。此版本要求 `platforms` 属性不为空。自 2023 年 6 月 7 日由 XlaCallModule 支持，自 2023 年 6 月 13 日（JAX 0.4.13）起支持 JAX。 + 第 7 版增加了对 `stablehlo.shape_assertion` 操作和在 `disabled_checks` 中指定的 `shape_assertions` 的支持。参见[形状多态性存在错误](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#errors-in-presence-of-shape-polymorphism)。自 2023 年 7 月 12 日（cl/547482522）由 XlaCallModule 支持，自 2023 年 7 月 20 日（JAX 0.4.14）起支持 JAX 序列化，并自 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 版添加了对 effects 的支持。详见 `export.Exported` 的文档字符串获取准确的调用约定。在此调用约定版本中，我们还使用 `jax.global_constant` 属性标记平台索引和维度变量参数。自 2023 年 10 月 27 日由 XlaCallModule 支持，自 2023 年 10 月 20 日（JAX 0.4.20）起支持 JAX，并自 2024 年 2 月 1 日（JAX 0.4.24）起成为默认选项。截至 2024 年 3 月 27 日，这是唯一支持的版本。 ## 从 `jax.experimental.export` 迁移指南。 在 2024 年 6 月 14 日，我们废弃了 `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` 对象的一个方法。不再使用 `export.call(exp)`，应使用 `exp.call`。 + `jax.experimental.export.serialize` 现在是 `jax.export.Exported` 对象的一个方法。不再使用 `export.serialize(exp)`，应使用 `exp.serialize()`。 + 配置标志 `--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`。
