<!--* freshness: { owner: "necula" reviewed: "2024-06-26" } *-->

# Exporting and serializing staged-out computations

The {ref}`ahead-of-time-lowering` APIs produce
objects that can be used for debugging or for compilation and
execution in the same process.
Sometimes you want to serialize a lowered JAX function for
compilation and execution in a separate process, perhaps
at a later time. This would allow you to:

  * compile and execute the function in another process or machine
    without requiring access to the JAX program,
    and without having to repeat the staging-out and lowering, e.g.,
    in an inference system.
  * trace and lower a function on a machine that does not have access
    to the accelerator for which you want to later compile and execute
    the function.
  * archive a snapshot of a JAX function, e.g., to be able to
    reproduce later your results. **Note:** check out the [compatibility
    guarantees](#compatibility-guarantees) for this use case.

For more details see the {mod}`jax.export` API reference.

Here is an example:

```python
>>> 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)

```

Serialization is broken down into two stages:
   1. exporting to produce an {class}`jax.export.Exported` object that contains
     the StableHLO for the lowered function along with the metadata necessary to
     call it from another JAX function. We have plans to add code to generate
     `Exported` objects from TensorFlow, and to use `Exported` objects from
     TensorFlow and PyTorch.
   2. the actual serialization to a byte array using the flatbuffers format.
     See {ref}`jax2tf` for
     an alternative serialization to TensorFlow graph that can be used
     for interoperation with TensorFlow.

**WARNING**: The serialized bytearray must be trusted input. If you execute it
after deserialization, it may execute any custom call registered in the jaxlib.

## Support for reverse-mode AD

Serialization can optionally support higher-order reverse-mode AD. This is done
by serializing the {func}`jax.vjp` of the primal function along with the primal
function,
up to a user-specified order (default is 0, meaning that the rehydrated
function cannot be differentiated):

```python
>>> 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)  # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
ValueError: No VJP is available

```

Note that the VJP function is computed lazily while serializing,
when the JAX program is still available.
This means that it respects all features of JAX VJP,
e.g., {func}`jax.custom_vjp` and {func}`jax.remat`.

Note that the rehydrated function does not support any other
transformations, e.g., forward-mode AD (jvp), or {func}`jax.vmap`.

## Compatibility guarantees

When you use the {class}`jax.export` module, you get the following export
compatibility guarantees:
A JAX exported and serialized artifact supports `.deserialize` and `.call`
(i.e. compilation and execution) by a compiler and JAX runtime system that are:

  * **up to 6 months newer** than the version of JAX used for exporting
  (we say that JAX export offers **6 months backward compatibility**).
  This is useful if you want to archive the exported artifact to be compiled and
  executed later.
  * **up to 3 weeks older** than the version of JAX used for exporting
  (we say that JAX export offers **3 weeks forward compatibility**).
  This is useful if you want to compile and run an exported artifact with a
  consumer that was built and deployed before the export, e.g.,
  an inference system that is already deployed when the exporting is done.

(The particular compatibility window lengths are the same that JAX
[promises for jax2tf](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#usage-saved-model),
and are based on [TensorFlow Compatibility](https://www.tensorflow.org/guide/versions#graph_and_checkpoint_compatibility_when_extending_tensorflow).
The terminology “backward compatibility” is from the perspective of the
consumer, e.g., the inference system.)

What **matters is when the exporting and consuming components were built**,
not the time when the exporting and the compilation happen.
To reduce chances of incompatibility, internal JAX users should **rebuild and
redeploy consumer systems as frequently as possible**.

External users should export for archival **with the latest released version of
jaxlib**.

The compatibility guarantees do not apply if you bypass the `jax.export` APIs
to obtain the StableHLO code (e.g., by `jax.jit(f).lower(1.).compiler_ir()`)
and then try to use it with a different build of JAX or jaxlib.
Unlike direct lowering, the {class}`jax.export` module uses the
[portable-artifact feature of StableHLO](https://github.com/openxla/stablehlo/blob/main/docs/compatibility.md)
to deal with the possible evolution of the StableHLO opset.

If you try to deserialize a model that is older than the backwards-compatibility
window you may get an error from the JAX deserializer, or may simply get
obscure errors or crashes from JAX or XLA. If you get a deserialization
error you can try to use the `--jax_export_deserialize_expired_versions=1`
flag (or the `JAX_EXPORT_DESERIALIZE_EXPIRED_VERSIONS=1` environment variable).
This may give you a small time window to continue your work while you work
to refresh the old serialized artifact.

**WARNING**: If you enable `--jax_export_deserialize_expired_versions=1` you
accept that the deserialization behavior may change at any time. You take
responsibility for refreshing the old serialized artifact ASAP.

### Compatibility guarantees for custom calls

Furthermore, the raw StableHLO may contain custom calls referencing C++
functions.
JAX uses custom calls for lowering of a small number of primitives,
e.g., linear algebra primitives, sharding annotations, or Pallas kernels.
These do not fall under the compatibility guarantees for StableHLO.
The C++ implementations of these functions change rarely, but they can change.

In order to ensure forward compatibility, when we change the JAX lowering rules
to use a new custom call target, JAX will refrain for 3 weeks to use the new
target. To use the latest lowering rules, you can pass the
`--jax_export_ignore_forward_compatibility=1` configuration flag
or the `JAX_EXPORT_IGNORE_FORWARD_COMPATIBILITY=1` environment variable.

Only a subset of custom calls are guaranteed stable and have
compatibility guarantees ([see list](https://github.com/search?q=repo%3Ajax-ml%2Fjax++%22_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE+%3D%22+path%3A_export.py&amp%3Btype=code&type=code)).
We continuously
add more custom call targets to the allowed list along with backwards
compatibility tests. If you try to serialize
code that invokes other custom call targets you will get an error
during exporting.

If you want to disable this safety check for a specific custom call,
e.g., with target `my_target`, you can add
`export.DisabledSafetyCheck.custom_call("my_target")` to the
`disabled_checks` parameter of the `export` method.

Alternatively, `export.DisabledSafetyCheck.custom_call("ALL")` can be used to
disable the check for all custom calls, this should only be used when export
compatibility is not needed.

The following example shows how to disable the check for a specific custom call:

```python
>>> 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.)  # doctest: +IGNORE_EXCEPTION_DETAIL
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.)

```

See {ref}`export_ensuring_compat` for developer information regarding
ensuring compatibility.

## Cross-platform and multi-platform export

JAX lowering is platform specific for a small number of JAX primitives.
By default, the code is lowered and exported for the accelerator
present on the exporting machine:

```python
>>> from jax import export
>>> export.default_export_platform()
'cpu'

```

There is a safety check that will raise an error when trying to compile
an `Exported` object on a machine that does not have the accelerator
for which the code was exported.

You can specify explicitly for what platforms the code should be exported.
This allows you to specify a different accelerator than you have
available at export time,
and it even allows you to specify multi-platform export to
obtain an `Exported` object that can be compiled and executed
on multiple platforms.

```python
>>> 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.)  # doctest: +IGNORE_EXCEPTION_DETAIL
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)

```

For multi-platform export, the StableHLO will contain multiple
lowerings but only for those primitives that require it, so the
resulting module size should be only marginally larger than the
size of a module with default export.
As an extreme case, when serializing a module without any
primitives with platform-specific lowering, you will get
the same StableHLO as for the single-platform export.

```python
>>> 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)  # doctest: +SKIP
9220

>>> exp_multi = export.export(jax.jit(f),
...                           platforms=["cpu", "tpu", "cuda"])(1.)
>>> len(exp_multi.mlir_module_serialized)  # doctest: +SKIP
9282

```

## Shape polymorphic export

When used in JIT mode, JAX will trace and lower a function separately
for each combination of input shapes. When exporting, it is possible
in some cases to use dimension variables for some input dimensions
in order to obtain an exported artifact that can be used with multiple
combinations of input shapes.

See the {ref}`shape_poly` documentation.

## Device-polymorphic export

An exported artifact may contain sharding annotations for inputs,
outputs and for some intermediates, but these annotations do not refer
directly to the actual physical devices that existed at exporting time.
Instead, the sharding annotations refer to logical devices. This
means that you can compile and run the exported artifacts on different
physical devices that were used for exporting.

The cleanest way to achieve a device-polymorphic export is to
use shardings constructed with a `jax.sharding.AbstractMesh`,
which contains only the mesh shape and axis names. But,
you can achieve the same results if you use shardings
constructed for a mesh with concrete devices, since the actual
devices in the mesh are ignored for tracing and lowering:

```python
>>> 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=(Explicit,)), 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),)']

```

It is an error to try to invoke an exported artifact with a different number
of devices than it was exported for:

```python
>>> 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)  # doctest: +IGNORE_EXCEPTION_DETAIL
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.

```

There are helper functions to shard the inputs for calling an exported
artifacts using a new mesh constructed at the call site:

```python
>>> 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)

```

As a special facility, if a function was exported for 1 device and if it
contains no sharding annotations, then it can be invoked on an argument of
the same shape but sharded on multiple devices, and the compiler will shard
the function appropriately:

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

```

## Calling convention versions

The JAX export support has evolved over time, e.g., to support
effects. In order to support compatibility (see [compatibility guarantees](#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](#calling-convention-versions)):

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

```

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:

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

```

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

(export-calling-convention-version)=

### 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).
  * Version 6 adds support for the `disabled_checks` attribute. This version
    mandates a non-empty `platforms` attribute. Supported by XlaCallModule
    since June 7th, 2023 and available in JAX since
    June 13th, 2023 (JAX 0.4.13).
  * Version 7 adds support for `stablehlo.shape_assertion` operations and
    for `shape_assertions` specified in `disabled_checks`.
    See [Errors in presence of shape polymorphism](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#errors-in-presence-of-shape-polymorphism).
    Supported by XlaCallModule
    since July 12th, 2023 (cl/547482522),
    available in JAX serialization since July 20th, 2023 (JAX 0.4.14),
    and the default since August 12th, 2023 (JAX 0.4.15).
  * Version 8 adds support for the `jax.uses_shape_polymorphism` module
    attribute and enables the shape refinement pass only when the
    attribute is present. Supported by XlaCallModule since July 21st, 2023
    (cl/549973693), available in JAX since July 26th, 2023 (JAX 0.4.14),
    and the default since October 21st, 2023 (JAX 0.4.20).
  * Version 9 adds support for effects.
    See the docstring for `export.Exported` for the precise calling convention.
    In this calling convention version we also tag the platform index and the
    dimension variables arguments with `jax.global_constant` attributes.
    Supported by XlaCallModule since October 27th, 2023,
    available in JAX since October 20th, 2023 (JAX 0.4.20),
    and the default since February 1st, 2024 (JAX 0.4.24).
    This is the only supported version as of 27th of March, 2024.
  * Version 10 propagate the `jax.config.use_shardy_partitioner` value to
    XlaCallModule. Supported by XlaCallModule since May 20th, 2025, and
    the default in JAX since July 14th, 2025 (JAX 0.7.0).

## Developer documentation

(export_debugging)=
### Debugging

You can log the exported modules, with somewhat different flags in OSS versus
in Google. In OSS you can do the following:

```shell
# 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
```

You will see a log line of the form:
```shell
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.
```

If you set the environment variable `JAX_DUMP_IR_TO` to a directory,
the exported (and the JIT compiled) HLO
modules will be saved there.

```shell
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.
```

You will see both the exported modules (named `..._export.mlir`
and the JIT compiled modules (named `..._compile.mlir`):
```shell
$ 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
```

Set [`JAX_DEBUG_LOG_MODULES=jax._src.export`](https://docs.jax.dev/en/latest/config_options.html#jax_debug_log_modules) to enable extra debugging logging.

(export_ensuring_compat)=
### Ensuring forward and backward compatibility

This section discusses the process JAX developers
should use to ensure the [compatibility guarantees](#compatibility-guarantees).

One complication is that external users install JAX and jaxlib
in separate packages,
and users often end up using an older jaxlib than JAX.
We observe that the custom calls live in the jaxlib, and only the jaxlib is
relevant for a consumer of an exported artifact.
To simplify the process, we are setting the expectation for external users
that the compatibility window is defined in terms of jaxlib releases,
and it is their responsibility to ensure that they export with a new jaxlib
even if JAX would function with an older version.

Thus, we care only about jaxlib releases.
We can start a backward-compatibility deprecation clock when we make a
jaxlib release, even if we don’t force it to be the minimum allowed version.

Let’s say that we need to add, delete, or change the semantics of a
custom call target `T` used by the JAX lowering rules.
Here is a possible chronology (for changing custom call targets
that live in jaxlib):

  1. Day “D - 1”, before the change. Say that the active internal JAX
     version is `0.4.31` (the version of the next JAX and jaxlib releases).
     The JAX lowering rules use a custom call `T`.
  2. Day “D”, we add the new custom call target `T_NEW`.
    We should create a new custom call target, and clean up the old
    target roughly after 6 months, rather than updating `T` in place:
       * See the example [PR #20997](https://github.com/jax-ml/jax/pull/20997)
         implementing the steps below.
       * We add the custom call target `T_NEW`.
       * We change the JAX lowering rules that were previous using `T`,
         to use `T_NEW`, conditionally as follows:

        ```python
        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", ...)
        ```
       * Note that the forward compatibility mode is always false in JIT mode
         or if the user passes `--jax_export_ignore_forward_compatibility=true`
       * Note that at this point the exports will still not use `T_NEW`.
  3. This can be done at any time after the previous step, and before
     the next step: Add a backward compatibility test for `T_NEW`,
     and add `T_NEW` to the list of
     [`_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE`](https://github.com/search?q=repo%3Ajax-ml%2Fjax++%22_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE+%3D%22+path%3A_export.py&amp%3Btype=code&type=code) in `_export.py`.
       * Instructions for adding backwards compatibility tests are at
         the top of
         [export_back_compat_test_util.py](https://github.com/search?q=repo%3Ajax-ml%2Fjax+++path%3Aexport_back_compat_test_util.py&amp%3Btype=code&type=code).
       * An example is in [PR #29488](https://github.com/jax-ml/jax/pull/29488).
       * Note that if you do this before the next step, the exporting will
         still not use the `T_NEW` lowering, and you have to add
         `with config.export_ignore_forward_compatibility(True):` around the
         call to `self.run_one_test`. This can be removed when you actually get
         to step 4.
       * You may also need to enable the test only for new versions of jaxlib.
  4. Day “D + 21” (end of forward compatibility window; can be even later than
    21 days):
    We remove the `forward_compat_mode` in the lowering code, so now exporting
    will start using the new custom call target `T_NEW` as long as we are using
    a new `jaxlib`.
  5. Day "RELEASE > D" (the first JAX release date after `D`, when we release
    version `0.4.31`):
    we start the clock for the 6 months backwards compatibility.
    Note that this is relevant only if `T` is among the custom call targets for
    which we already guarantee stability, i.e., are listed in
    [`_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE`](https://github.com/search?q=repo%3Ajax-ml%2Fjax++%22_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE+%3D%22+path%3A_export.py&amp%3Btype=code&type=code).
      * If `RELEASE` is in the forward compatibility window `[D, D + 21]` and if
        we make `RELEASE` the minimum allowed jaxlib version then we can
        remove the `jaxlib_version < (0, 4, 31)` conditional in the
        JIT branch.
  6. Day “RELEASE + 180” (end of backward compatibility window,
    can be even later than 180 days): By now, we must have bumped
    the minimum jaxlib so that the lowering conditional
    `jaxlib_version < (0, 4, 31)`
    was already removed and JAX lowering cannot generate custom calls to `T`.
      * We remove the C++ implementation of the old custom call target `T`.
      * We remove also the backwards compatibility test for `T`

## Migration guide from jax.experimental.export

On June 18, 2024 (JAX version 0.4.30)
we deprecated the `jax.experimental.export` APIs
in favor of `jax.export` APIs. There have been some minor changes:

  * `jax.experimental.export.export`:
    * The old function used to allow any Python callable, or the result of
      `jax.jit`. Now only the latter is accepted. You have to manually apply
      `jax.jit` to the function to export before calling `export`.
    * The old `lowering_parameters` kwarg is now named `platforms`
  * `jax.experimental.export.default_lowering_platform()` is now
    at {func}`jax.export.default_export_platform`.
  * `jax.experimental.export.call` is now a method of the
    {class}`jax.export.Exported` object.
    Instead of `export.call(exp)` you should use `exp.call`.
  * `jax.experimental.export.serialize` is now a method of the
    {class}`jax.export.Exported`
    object. Instead of `export.serialize(exp)` you should use `exp.serialize()`.
  * The configuration flag `--jax-serialization-version` is deprecated.
    Use `--jax-export-calling-convention-version`.
  * The value `jax.experimental.export.minimum_supported_serialization_version`
    is now at `jax.export.minimum_supported_calling_convention_version`.
  * The following fields of {class}`jax.export.Exported` have been renamed
     * `uses_shape_polymorphism` is now `uses_global_constants`
     * `mlir_module_serialization_version` is now `calling_convention_version`
     * `lowering_platforms` is now `platforms`.
