export.py

# Copyright 2023 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""JAX APIs for exporting JAX functions for interoperation.

"""

from collections.abc import Sequence
import copy
import dataclasses
import functools
import itertools
import re
from typing import Any, Callable, Optional, Union

from absl import logging

import numpy as np

import jax
from jax import config
from jax import sharding

from jax._src import core
from jax._src import dispatch
from jax._src.interpreters import mlir
from jax._src.interpreters import pxla
from jax._src.lib import xla_client
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import hlo
from jax._src.lib.mlir.dialects import func as func_dialect
from jax._src import pjit
from jax._src import sharding_impls
from jax._src import source_info_util
from jax._src import tree_util
from jax._src import util
from jax._src import xla_bridge as xb

from jax.experimental.export import shape_poly

map = util.safe_map
zip = util.safe_zip

DType = Any

class DisabledSafetyCheck:
  """A safety check should be skipped on (de)serialization.

  Most of these checks are performed on serialization, but some are deferred to
  deserialization. The list of disabled checks is attached to the serialization,
  e.g., as a sequence of string attributes to `jax_export.Exported` or of
  `tf.XlaCallModuleOp`.

  You can disable more deserialization safety checks by passing
  `TF_XLA_FLAGS=--tf_xla_call_module_disabled_checks=platform`.
  """
  _impl: str

  @classmethod
  def platform(cls) -> "DisabledSafetyCheck":
    """Allows the execution platform to differ from the serialization platform.

    Has effect only on deserialization.
    """
    return DisabledSafetyCheck("platform")

  @classmethod
  def custom_call(cls, target_name: str) -> "DisabledSafetyCheck":
    """Allows the serialization of a call target not known to be stable.

    Has effect only on serialization.
    Args:
      target_name: the name of the custom call target to allow.
    """
    return DisabledSafetyCheck(f"custom_call:{target_name}")

  @classmethod
  def shape_assertions(cls) -> "DisabledSafetyCheck":
    """Allows invocations with shapes that do not meet the constraints.

    Has effect on serialization (to suppress the generation of the assertions)
    and also on deserialization (to suppress the checking of the assertions).
    """
    return DisabledSafetyCheck("shape_assertions")

  def is_custom_call(self) -> Optional[str]:
    """Returns the custom call target allowed by this directive."""
    m = re.match(r'custom_call:(.+)$', self._impl)
    return m.group(1) if m else None

  def __init__(self, _impl:str):
    # Do not use directly, use builders `platform`, `custom_call`.
    self._impl = _impl

  def __str__(self):
    return self._impl
  __repr__ = __str__

  def __eq__(self, other) -> bool:
    return isinstance(other, DisabledSafetyCheck) and self._impl == other._impl

  def __hash__(self) -> int:
    return hash(self._impl)


minimum_supported_serialization_version = 6
maximum_supported_serialization_version = 8

@dataclasses.dataclass(frozen=True)
class Exported:
  """A JAX function lowered to StableHLO.

  Attributes:
    fun_name: the name of the exported function, for error messages.
    in_tree: a PyTreeDef describing the tuple (args, kwargs) of the lowered JAX
        function. The actual lowering does not depend on the `in_tree`, but this
        can be used to invoke the exported function using the same argument
        structure.
    in_avals: the flat tuple of input abstract values. May contain dimension
        expressions in the shapes.
    out_tree: a PyTreeDef describing the result of the lowered JAX function.
    out_avals: the flat tuple of output abstract values. May contain dimension
        expressions in the shapes, with dimension variables among those in
        `in_avals.
    in_shardings: the flattened input shardings. Only for the inputs that are
        specified in `module_kept_var_idx`. If `None` then it is equivalent
        to unspecified shardings.
    out_shardings: the flattened output shardings, as long as `in_avals`.
    lowering_platforms: a tuple containing at least one of 'tpu', 'cpu',
        'cuda', 'rocm'. See below for the calling convention for when
        there are multiple lowering platforms.
    mlir_module_serialized: the serialized lowered VHLO module.
    serialization_version: a version number for the serialized module.
        See more versioning details at https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions.
    module_kept_var_idx: the sorted indices of the arguments among `in_avals` that
        must be passed to the module. The other arguments have been dropped
        because they are not used. Same length as `in_shardings`.
    uses_shape_polymorphism: whether the `mlir_module_serialized` uses shape
        polymorphism. This may be because `in_avals` contains dimension
        variables, but also from inner calls of shape-polymorphic
        Exported modules.
    disabled_checks: a list of descriptors of safety checks that have been
        disabled at export time. See docstring for `DisabledSafetyCheck`.
    _get_vjp: an optional function that takes the current exported function and
        returns the exported VJP function.
        The VJP function takes a flat list of arguments,
        starting with the primal arguments and followed by a cotangent argument
        for each primal output. It returns a tuple with the cotangents
        corresponding to the flattened primal inputs.

  Calling convention for the exported module:

  The `mlir_module` has a `main` function that takes an optional first
  platform index argument if the module supports multiple platforms
  (`len(lowering_platforms) > 1`), followed by the kept array arguments
  (corresponding to `module_kept_var_idx` and `in_avals`).
  The platform index is a i32 scalar encoding the index of the current
  compilation platform into the `lowering_platforms` sequence.

  Inner functions use a different calling convention: an optional
  platform index argument, optional dimension variable arguments specified
  using scalar tensors of type i32 or i64,
  followed by optional token arguments (in presence of side 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
  ordered effects. The `main` function will be as follows:

      func public main(platform_index: i32, 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 res
      }

  The actual computation is in `_wrapped_jax_export_main`, taking also
  the values of `h` and `w` and the token. Proper exporting of
  functions with side-effects and tokens is still work-in-progress.

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

  If we `call_exported` with this module we perform these checks
  statically (in `call_exported_abstract_eval`).
  """
  fun_name: str
  in_tree: tree_util.PyTreeDef
  in_avals: tuple[core.AbstractValue, ...]
  out_tree: tree_util.PyTreeDef
  out_avals: tuple[core.AbstractValue, ...]

  in_shardings: Optional[tuple[Union[sharding.XLACompatibleSharding, pxla.UnspecifiedValue], ...]]
  out_shardings: Optional[tuple[Union[sharding.XLACompatibleSharding, pxla.UnspecifiedValue], ...]]
  lowering_platform: str  # For backwards compatibility
  lowering_platforms: tuple[str, ...]
  disabled_checks: Sequence[DisabledSafetyCheck]

  mlir_module_serialized: bytes
  serialization_version: int
  module_kept_var_idx: tuple[int, ...]
  uses_shape_polymorphism: bool

  _get_vjp: Optional[Callable[["Exported"], "Exported"]]

  def mlir_module(self) -> ir.Module:
    return xla_client._xla.mlir.deserialize_portable_artifact(self.mlir_module_serialized)

  def __str__(self):
    # This is called to make a MLIR source location when we call an Exported, and we
    # do not want the entire serialized module to end up in locations.
    return f"Exported(fun_name={self.fun_name}, ...)"

  def vjp(self) -> "Exported":
    """Gets the exported VJP.

    Returns None if not available, which can happen if the Exported has been
    loaded from an external format, without a VJP."""
    if self._get_vjp is None:
      raise ValueError("No VJP is available")
    return self._get_vjp(self)


def default_lowering_platform() -> str:
  # Canonicalize to turn 'gpu' into 'cuda' or 'rocm'
  return xb.canonicalize_platform(jax.default_backend())

def poly_spec(
    arg_shape: Sequence[Optional[int]],
    arg_dtype: DType,
    polymorphic_shape: Optional[str]) -> jax.ShapeDtypeStruct:
  """Constructs a jax.ShapeDtypeStruct with polymorphic shapes.

  Args:
    arg_shape: the shape, with possibly some unspecified dimensions.
    arg_dtype: the jax dtype.
    polymorphic_shape: a string specifying the polymorphic shape.

      .. warning:: The shape-polymorphic lowering is an experimental feature.
        It is meant to be sound, but it is known to reject some JAX programs
        that are shape polymorphic. The details of this feature can change.

      It should be either `None` (all dimensions are constant), or a string of
      specification for one axis, and can be either a constant, `_` denoting
      a constant dimension given by the `arg_shape`, or the name of a
      dimension variable assumed to range over dimension greater than 0. For
      convenience, zero or more trailing `_` can be abbreviated with `...`, and
      the surrounding parentheses may be missing.

      Note that this function does not ensure that the provided `arg_shape`
      is compatible with `polymorphic_shape`. The `arg_shape` is used only
      to fill-in placeholders from `polymorphic_shape`.

      See [the README](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#shape-polymorphic-conversion)
      for more details.

  Returns: a jax.ShapeDTypeStruct with shapes that may contain symbolic
      expressions involving dimension variables.
  """
  aval_shape = shape_poly._parse_spec(polymorphic_shape, arg_shape)
  return jax.ShapeDtypeStruct(aval_shape, arg_dtype)

def shape_and_dtype_jax_array(a) -> tuple[Sequence[Optional[int]], DType]:
  """Returns the shape and dtype of a jax.Array."""
  aval = core.raise_to_shaped(core.get_aval(a))
  return aval.shape, aval.dtype

def poly_specs(
    args,  # pytree of arguments
    polymorphic_shapes,  # prefix pytree of strings
    get_shape_and_dtype=shape_and_dtype_jax_array,
):
  """Constructs a pytree of jax.ShapeDtypeSpec.

  Args:
    args: a pytree of arguments
    polymorphic_shapes: should be `None` (all arguments are monomorphic),
      a single string (applies to all arguments), or a pytree matching a prefix
      of the `args`.
      See [how optional parameters are matched to
      arguments](https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees).

      Note that this function does not ensure that the provided `args` shapes
      are compatible with `polymorphic_shapes`. The `args.shape` are used only
      to fill-in placeholders from `polymorphic_shapes`.

      See docstring of `poly_spec` and
      [the README](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#shape-polymorphic-conversion)
      for more details.

  Returns: a pytree of jax.ShapeDTypeStruct matching `args`.
  """
  args_flat, args_tree = tree_util.tree_flatten(args)

  shapes_and_dtypes = tuple(map(get_shape_and_dtype, args_flat))
  shapes, dtypes = util.unzip2(shapes_and_dtypes)

  if isinstance(args, tuple) and isinstance(polymorphic_shapes, list):
    # TODO: Remove backward-compatibility workaround
    polymorphic_shapes_ = tuple(polymorphic_shapes)
  else:
    polymorphic_shapes_ = polymorphic_shapes

  try:
    polymorphic_shapes_flat = tree_util.broadcast_prefix(
        polymorphic_shapes_, args,
        is_leaf=lambda x: x is None)
  except ValueError:
    e, *_ = tree_util.prefix_errors(
        polymorphic_shapes_, args,
        is_leaf=lambda x: x is None)
    raise e("jax_export polymorphic_shapes") from None

  # Now add in the polymorphic shapes
  args_specs_flat = tuple(
      map(poly_spec, shapes, dtypes, polymorphic_shapes_flat))

  return args_tree.unflatten(args_specs_flat)


def export(fun_jax: Callable,
           *,
           lowering_platform: Optional[str] = None,
           lowering_platforms: Optional[Sequence[str]] = None,
           disabled_checks: Sequence[DisabledSafetyCheck] = (),
           ) -> Callable[..., Exported]:
  """Exports native serialization for a JAX function.

  Args:
    fun_jax: the function to lower and serialize.
    lowering_platform: one of 'tpu', 'cpu', 'cuda', 'rocm'. If None, then use
        the default JAX backend.
    lowering_platforms: DO NOT USE (NOT YET FUNCTIONAL).
        Optional sequence containing a subset of 'tpu', 'cpu',
        'cuda', 'rocm'. If more than one platform is specified, then
        the lowered code takes an argument specifying the platform.
        If None, then use the default JAX backend.
        The calling convention for multiple platforms is explained in the
        `jax_export.Exported` docstring.
    disabled_checks: the safety checks to disable. See docstring
        of `DisabledSafetyCheck`.

  Returns: a function that takes args and kwargs pytrees of jax.ShapeDtypeStruct,
      or values with `.shape` and `.dtype` attributes, and returns an
      `Exported`.

  Usage:

      def f_jax(*args, **kwargs): ...
      exported = jax_export.export(f_jax)(*args, **kwargs)
  """
  fun_name = getattr(fun_jax, "__name__", "unknown")
  version = config.jax_serialization_version
  if (version < minimum_supported_serialization_version or
      version > maximum_supported_serialization_version):
    raise ValueError(
      f"The requested jax_serialization version {version} is outside the "
      f"range of supported versions [{minimum_supported_serialization_version}"
      f"..{maximum_supported_serialization_version}]")

  def do_export(*args_specs, **kwargs_specs) -> Exported:
    if not hasattr(fun_jax, "lower"):
      # We support convert(pjit(f_jax)) and convert(jit(f_jax)) but also
      # convert(f_jax), in which case a "jit" is implied. In that case we raise
      # an error if the lowered function contains non-replicated sharding annotations.
      wrapped_fun_jax = jax.jit(fun_jax)
      allow_non_replicated_sharding = False
    else:
      # If we have a pjit or pmap already we do not wrap with another, and we
      # allow shardings.
      wrapped_fun_jax = fun_jax  # type: ignore
      allow_non_replicated_sharding = True

    nonlocal lowering_platforms
    if lowering_platforms is not None:
      lowering_platforms = tuple(lowering_platforms)
    else:
      lowering_platforms = (lowering_platform or default_lowering_platform(),)

    # Do not include shape assertions if the version is < 7.
    enable_shape_assertions = (
        DisabledSafetyCheck.shape_assertions() not in disabled_checks and
        version >= 7)  # type: ignore
    try:
      prev_enable_shape_assertions = shape_poly.thread_local_state.enable_shape_assertions
      shape_poly.thread_local_state.enable_shape_assertions = enable_shape_assertions
      lowered = wrapped_fun_jax.lower(
          *args_specs, **kwargs_specs,
          _experimental_lowering_platform=lowering_platforms)

      lowering = lowered._lowering  # type: ignore
      _check_lowering(lowering)
      mlir_module = lowering.stablehlo()

      args_avals_flat, _ = tree_util.tree_flatten(lowered.in_avals)
      if "kept_var_idx" in lowering.compile_args:
        module_kept_var_idx = tuple(sorted(lowering.compile_args["kept_var_idx"]))
      else:
        # For pmap
        module_kept_var_idx = tuple(range(len(args_avals_flat)))
      shape_poly_state = lowering.compile_args["shape_poly_state"]
      if (not all(core.is_constant_shape(a.shape) for a in args_avals_flat)
          or lowering.compile_args.get("ordered_effects", [])):
        mlir_module = _wrap_main_func(
            mlir_module, args_avals_flat, args_kwargs_tree=lowered.in_tree,
            has_platform_index_argument=shape_poly_state.has_platform_index_argument,
            module_kept_var_idx=module_kept_var_idx)
    finally:
      shape_poly.thread_local_state.enable_shape_assertions = prev_enable_shape_assertions

    with mlir_module.context:
      mlir_module_attrs = mlir_module.operation.attributes
      mlir_module_attrs["jax.uses_shape_polymorphism"] = (
          mlir.ir.BoolAttr.get(shape_poly_state.uses_dim_vars))

    mlir_module_serialized = _serialize_module(mlir_module)

    # Figure out the result types and shapes
    if "global_out_avals" in lowering.compile_args:
      # This is currently the case for pjit
      out_avals_flat = lowering.compile_args["global_out_avals"]
    elif "shards" in lowering.compile_args:  # for PmapComputation
      out_avals_flat = lowering.compile_args["shards"].out_sharded_avals
    else:
      out_avals_flat = lowered.compile_args["out_avals"]

    # Log and then check the module.
    if logging.vlog_is_on(3):
      mlir_module_text = mlir.module_to_string(mlir_module)
      logmsg = (f"version={version} "
                f"lowering_platforms={lowering_platforms} "
                f"disabled_checks={disabled_checks}")
      logging.info("Lowered JAX module: %s\n", logmsg)
      for l in mlir_module_text.splitlines():
        logging.info(l)

    _check_module(mlir_module,
                  allow_non_replicated_sharding=allow_non_replicated_sharding,
                  disabled_checks=disabled_checks)

    return Exported(
        fun_name=fun_name,
        in_tree=lowered.in_tree,
        out_tree=lowered.out_tree,
        in_avals=tuple(args_avals_flat),
        out_avals=tuple(out_avals_flat),
        in_shardings=lowering.compile_args["in_shardings"],
        out_shardings=lowering.compile_args["out_shardings"],
        lowering_platform=lowering_platforms[0],  # TODO: remove
        lowering_platforms=lowering_platforms,
        disabled_checks=tuple(disabled_checks),
        mlir_module_serialized=mlir_module_serialized,
        module_kept_var_idx=module_kept_var_idx,
        uses_shape_polymorphism=shape_poly_state.uses_dim_vars,
        serialization_version=version,  # type: ignore
        _get_vjp=lambda exported: _export_native_vjp(fun_jax, exported))

  return do_export


def _serialize_module(module: ir.Module) -> bytes:
  mlir_str = mlir.module_to_bytecode(module)
  if hlo.get_api_version() < 4:
    target_version = hlo.get_earliest_forward_compatible_version()
  else:
    # `target_version` is used to manage situations when a StableHLO producer
    # (in this case, jax2tf) and a StableHLO consumer were built using
    # different versions of StableHLO.
    #
    # Each StableHLO version `producer_version` has a compatibility window,
    # i.e. range of versions [`consumer_version_min`, `consumer_version_max`],
    # where StableHLO portable artifacts serialized by `producer_version`
    # can be deserialized by `consumer_version` within the window.
    # See https://github.com/openxla/stablehlo/blob/main/docs/compatibility.md
    # for the exact extent of these compatibility guarantees.
    #
    # `hlo.get_minimum_version()` returns `consumer_version_min`
    # for the current version of StableHLO. We are using it here to maximize
    # forward compatibility, i.e. to maximize how far into the past we can go
    # and still have the payloads produced by `serialize_portable_artifact`
    # compatible with potential consumers from the past.
    target_version = hlo.get_minimum_version()
  module_serialized = xla_client._xla.mlir.serialize_portable_artifact(
      mlir_str, target_version)
  return module_serialized


def _wrap_main_func(
    module: ir.Module,
    args_avals_flat: Sequence[core.ShapedArray],
    *,
    args_kwargs_tree: tree_util.PyTreeDef,
    has_platform_index_argument: bool,
    module_kept_var_idx: tuple[int, ...]
) -> ir.Module:
  """Wraps the lowered module with a new "main" handling dimension arguments.

  See calling convention documentation for `jax_export.Exported`.

  Args:
    module: the HLO module as obtained from lowering. See the calling convention
      for inner functions in `jax_export.Exported`.
    args_avals_flat: the avals for all the arguments of the lowered function,
      which correspond to the array arguments of the `module`.
    args_kwargs_tree: the PyTreeDef corresponding to `(args, kwargs)`, for error
      messages.
    has_platform_index_argument: whether the `module` has a first platform
      index argument
    module_kept_var_idx: a sorted tuple of integers with the indices of arguments
      in `args_avals_flat` that are kept as `module` arguments.

  Returns the wrapped module, without dimension and token arguments.
  """
  dim_vars = shape_poly.all_dim_vars(args_avals_flat)
  context = mlir.make_ir_context()
  with context, ir.Location.unknown(context):
    # Make a copy, do not mutate because it may be cached
    wrapped_module = ir.Module.parse(mlir.module_to_bytecode(module))
    symbol_table = ir.SymbolTable(wrapped_module.operation)
    orig_main = symbol_table["main"]
    orig_main.attributes["sym_visibility"] = ir.StringAttr.get("private")
    symbol_table.set_symbol_name(orig_main, "_wrapped_jax_export_main")
    orig_main_name = ir.StringAttr(symbol_table.insert(orig_main)).value

    def is_token(attrs):
      try:
        return ir.BoolAttr(ir.DictAttr(attrs)["jax.token"]).value
      except KeyError:
        return False

    orig_input_types = orig_main.type.inputs
    arg_attrs = list(ir.ArrayAttr(orig_main.arg_attrs))
    # The order of args: platform_index_arg, dim args, token args, array args.
    nr_platform_index_args = 1 if has_platform_index_argument else 0
    nr_dim_args = len(dim_vars)
    nr_token_args = sum(1 for attrs in arg_attrs if is_token(attrs))
    nr_array_args = len(orig_input_types) - nr_platform_index_args - nr_dim_args - nr_token_args
    assert nr_array_args >= 0
    assert not any(is_token(attrs) for attrs in arg_attrs[-nr_array_args:])
    (platform_input_types, dim_var_input_types,
     token_input_types, array_input_types) = util.split_list(
      orig_input_types, [nr_platform_index_args, nr_dim_args, nr_token_args])
    new_main_input_types = platform_input_types + array_input_types
    orig_output_types = orig_main.type.results
    result_attrs = list(ir.ArrayAttr(orig_main.result_attrs))
    nr_token_results = sum(1 for attrs in result_attrs if is_token(attrs))
    nr_array_results = len(orig_output_types) - nr_token_results
    assert nr_array_results >= 0
    assert not any(
        is_token(attrs) for attrs in result_attrs[-nr_array_results:])
    new_main_output_types = orig_output_types[-nr_array_results:]
    new_main_ftype = ir.FunctionType.get(new_main_input_types, new_main_output_types)
    new_main_op = func_dialect.FuncOp(
        "main", new_main_ftype, ip=ir.InsertionPoint.at_block_begin(wrapped_module.body))
    new_main_op.attributes["sym_visibility"] = ir.StringAttr.get("public")
    try:
      new_main_op.arg_attrs = ir.ArrayAttr.get(arg_attrs[0:nr_platform_index_args] + arg_attrs[-nr_array_args:])
    except KeyError:
      pass  # TODO: better detection if orig_main.arg_attrs does not exist
    try:
      new_main_op.result_attrs = ir.ArrayAttr.get(
          result_attrs[-nr_array_results:])
    except KeyError:
      pass
    symbol_table.insert(new_main_op)
    entry_block = new_main_op.add_entry_block()
    with ir.InsertionPoint(entry_block):
      module_context = mlir.ModuleContext(
          "cpu", "cpu", sharding_impls.ShardingContext([]),
          source_info_util.new_name_stack(),
          [], itertools.count(1), [], module=wrapped_module, context=context)
      ctx = mlir.LoweringRuleContext(
        module_context=module_context, primitive=None,
        avals_in=args_avals_flat, avals_out=None,
        tokens_in=mlir.TokenSet(), tokens_out=None)
      # We compute dim_values from the array arguments.
      new_main_op_array_args = new_main_op.arguments[nr_platform_index_args:]
      if shape_poly.all_dim_vars(args_avals_flat):
        # TODO(necula): handle module_kept_var_idx in presence of shape
        # polymorphism. For now we ensured upstream that we keep all variables.
        assert len(set(module_kept_var_idx)) == len(args_avals_flat)
        dim_values = mlir.lower_fun(
            functools.partial(shape_poly.compute_dim_vars_from_arg_shapes,
                              args_avals_flat, args_kwargs_tree=args_kwargs_tree),
            multiple_results=True)(ctx, *new_main_op_array_args)
      else:
        dim_values = ()
      # The arguments to pass to the call to orig_main
      orig_main_args: list[ir.Value] = []
      # The platform index and the dimension variables
      for arg, arg_type in zip(
          list(new_main_op.arguments[0:nr_platform_index_args]) + util.flatten(dim_values),
          platform_input_types + dim_var_input_types):
        if arg.type != arg_type:
          orig_main_args.append(hlo.ConvertOp(arg_type, arg).result)
        else:
          orig_main_args.append(arg)
      # Then the token arguments
      orig_main_args.extend(list(mlir.dummy_token()) * nr_token_args)
      # Then the array arguments. We insert a ConvertOp as the only use of
      # an input argument. This helps the downstream shape refinement because
      # it will set the type of input arguments to static shapes, and this
      # can invalidate the module if the argument is used as the result of a
      # function, or if it appears as the input to a custom_call with
      # output_operand_alias attribute. See b/287386268.
      for a in new_main_op_array_args:
        orig_main_args.append(hlo.ConvertOp(a.type, a).result)
      call = func_dialect.CallOp(orig_output_types,
                                 ir.FlatSymbolRefAttr.get(orig_main_name),
                                 orig_main_args)
      func_dialect.ReturnOp(call.results[-nr_array_results:])
    symbol_table.set_symbol_name(new_main_op, "main")

  return wrapped_module

def _check_lowering(lowering) -> None:
  if not isinstance(lowering, pxla.MeshComputation):
    raise NotImplementedError(f"serialization is supported only for pjit. {lowering}")

  if lowering.compile_args["host_callbacks"] or lowering.compile_args["keepalive"]:
    raise NotImplementedError("serialization of host_callbacks is not yet implemented")
  # Check that we do not see new compile_args. When we add a compile_args it is
  # safe to add it to the allowed_compile_args if it does not change the semantics
  # or the calling convention of the lowered module.
  allowed_compile_args = [
      "backend", "mesh", "global_in_avals",
      "global_out_avals", "in_shardings", "out_shardings", "kept_var_idx",
      "spmd_lowering", "auto_spmd_lowering",
      "tuple_args", "ordered_effects", "unordered_effects",
      "keepalive", "host_callbacks", "pmap_nreps", "committed",
      "device_assignment", "jaxpr_debug_info", "shape_poly_state",
      "all_default_mem_kind"]
  for compile_arg in lowering.compile_args.keys():
    if compile_arg not in allowed_compile_args:
      raise NotImplementedError(f"Unrecognized lowered.compile_args[{compile_arg}]")

  # We have not implemented support for some of the compile_args. Check here that
  # the compile_args have the values that have been implemented.
  not_implemented_msgs = []
  for compile_arg, check_value, err_msg in (
      ("spmd_lowering", lambda v: v, "True"),
      ("auto_spmd_lowering", lambda v: not v, "False"),
      # tuple_args is a compilation flag, does not affect lowering.
      ("tuple_args", lambda v: True, "N/A"),
      # unordered_effects do not change the calling convention. Those from
      # jax.debug will also result in keepalive being non-empty and unsupported
      # custom calls. The CallTfEffect is an exception, but we want to allow
      # that one.
      ("unordered_effects", lambda v: True, "N/A"),
      # ordered_effects are allowed and we ensure that the calling convention is
      # unmodified by passing dummy tokens in the main function wrapper.
      ("ordered_effects", lambda v: True, "N/A"),
      # used for TPU jax.debug, send/recv. Not supported yet.
      ("host_callbacks", lambda v: not v, "empty"),
      # used on all platforms for callbacks. Not supported yet.
      ("keepalive", lambda v: not v, "empty"),
      ("pmap_nreps", lambda v: v == 1, "1"),
      ("shape_poly_state", lambda v: True, "N/A"),
  ):
    if compile_arg in lowering.compile_args:
      if not check_value(lowering.compile_args[compile_arg]):
        not_implemented_msgs.append(
            f"{compile_arg} must be {err_msg} and it is {lowering.compile_args[compile_arg]}")
  if not_implemented_msgs:
    raise NotImplementedError(
        "serialization error, unimplemented lowered.compile_args:\n" +
        "\n".join(not_implemented_msgs))

# These are the JAX custom call target names that are guaranteed to be stable.
# Their backwards compatibility is tested by back_compat_test.py.
_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE = {
    "Sharding", "SPMDFullToShardShape", "SPMDShardToFullShape",
    "ducc_fft", "dynamic_ducc_fft", "cu_threefry2x32",
    # cholesky on CPU
    "lapack_spotrf", "lapack_dpotrf", "lapack_cpotrf", "lapack_zpotrf",
    # eigh on CPU
    "lapack_ssyevd", "lapack_dsyevd", "lapack_cheevd", "lapack_zheevd",
    # eigh on GPU
    "cusolver_syevj", "cusolver_syevd",
    # eigh on TPU
    "Eigh",
    # eig on CPU
    "lapack_sgeev", "lapack_dgeev", "lapack_cgeev", "lapack_zgeev",
    # qr on CPU
    "lapack_sgeqrf", "lapack_dgeqrf", "lapack_cgeqrf", "lapack_zgeqrf",
    # householder product on CPU
    "lapack_sorgqr", "lapack_dorgqr", "lapack_cungqr", "lapack_zungqr",
    # svd on CPU
    "lapack_sgesdd", "lapack_dgesdd", "lapack_cgesdd", "lapack_zgesdd",
    # qr on GPU
    "cusolver_geqrf", "cublas_geqrf_batched",
    "cusolver_geqrf", "cusolver_orgqr",
    # qr and svd on TPU
    "Qr", "ProductOfElementaryHouseholderReflectors",
    # triangular_solve on CPU
    "blas_strsm", "blas_dtrsm", "blas_ctrsm", "blas_ztrsm",
    # TODO(atondwal, necula): add back_compat tests for lu on CPU/GPU
    # # lu on CPU
    "lapack_sgetrf",  "lapack_dgetrf", "lapack_cgetrf", "lapack_zgetrf",
    # schur on CPU
    "lapack_sgees", "lapack_dgees", "lapack_cgees", "lapack_zgees",
    # # lu on GPU
    # "cublas_getrf_batched", "cusolver_getrf",
    # "hipblas_getrf_batched", "hipsolver_getrf",
    # lu on TPU
    "LuDecomposition",
    # ApproxTopK on TPU
    "ApproxTopK",
    "tf.call_tf_function",  # From jax2tf.call_tf(func, call_tf_graph=True)
    "tpu_custom_call",  # Pallas/TPU kernels
    # TODO(burmako): maintain backwards compatibility for these, until they
    # are upstreamed to StableHLO.
    # See https://github.com/openxla/stablehlo/issues/8.
    "stablehlo.dynamic_reduce_window",
    "stablehlo.dynamic_rng_bit_generator",
    "stablehlo.dynamic_top_k",
    "shape_assertion",  # Used by shape_poly to evaluate assertions
}


def _check_module(mod: ir.Module, *,
                  allow_non_replicated_sharding: bool,
                  disabled_checks: Sequence[DisabledSafetyCheck]) -> None:
  """Run a number of checks on the module.

  Args:
    allow_non_replicated_sharding: whether the module is allowed to contain
      non_replicated sharding annotations.
    disabled_checks: the safety checks that are disabled.
  """
  sharding_attr = ir.StringAttr.get("Sharding", mod.context)
  shape_assertion_attr = ir.StringAttr.get("shape_assertion", mod.context)
  allowed_custom_call_targets: set[str] = copy.copy(_CUSTOM_CALL_TARGETS_GUARANTEED_STABLE)
  for dc in disabled_checks:
    target = dc.is_custom_call()
    if target is not None:
      allowed_custom_call_targets.add(target)

  allowed_custom_call_targets_attrs = {
      ir.StringAttr.get(target, mod.context)
      for target in allowed_custom_call_targets}
  disallowed_custom_call_ops: list[str] = []
  def check_sharding(op: ir.Operation, loc: ir.Location):
    if not allow_non_replicated_sharding:
      try:
        sharding = op.attributes["mhlo.sharding"]
      except KeyError:
        pass
      else:
        if ir.StringAttr(sharding).value not in ["{replicated}", ""]:
          raise ValueError(
              "Lowered function does not have a top-level pjit but it has"
              f" non-replicated sharding annotations, e.g., {op} at {loc}.\nSee"
              " https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#support-for-partitioning"
              " for a discussion."
          )

  def check_op(op: ir.Operation):
    op_name = op.operation.name
    if op_name == "func.func":
      check_sharding(op.operation, op.location)

    elif op_name == "stablehlo.custom_call" or op_name == "mhlo.custom_call":
      call_target_name_attr = op.operation.attributes["call_target_name"]
      if (call_target_name_attr not in allowed_custom_call_targets_attrs):
        disallowed_custom_call_ops.append(f"{op} at {op.location}")
      if call_target_name_attr == sharding_attr:
        check_sharding(op, op.location)
      elif call_target_name_attr == shape_assertion_attr:
        assert (DisabledSafetyCheck.shape_assertions() not in disabled_checks)

  def walk_operations(op):
    check_op(op)
    for region in op.operation.regions:
      for block in region:
        for op in block:
          walk_operations(op)

  walk_operations(mod)
  if disallowed_custom_call_ops:
    disallowed_custom_call_ops_str = "\n".join(disallowed_custom_call_ops)
    msg = ("Cannot serialize code with custom calls whose targets have no "
           "compatibility guarantees. Examples are:\n"
           f"{disallowed_custom_call_ops_str}.\n"
           "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-lowering-supports-only-select-custom-calls")
    raise ValueError(msg)


def _get_vjp_fun(primal_fun: Callable, *,
                 in_tree: tree_util.PyTreeDef,
                 in_avals: Sequence[core.AbstractValue],
                 out_avals: Sequence[core.AbstractValue],
                 module_kept_var_idx: tuple[int, ...],
                 in_shardings,
                 out_shardings,
                 apply_jit: bool
                 ) -> tuple[Callable, Sequence[core.AbstractValue]]:
  # Since jax.vjp does not handle kwargs, it is easier to do all the work
  # here with flattened functions.
  def fun_vjp_jax(*args_and_out_cts_flat_jax):
    # Takes a flat list of primals and output cotangents
    def flattened_primal_fun_jax(*args_flat):
      args, kwargs = in_tree.unflatten(args_flat)
      res = primal_fun(*args, **kwargs)
      res_flat, _ = tree_util.tree_flatten(res)
      return res_flat

    args_flat_jax, out_cts_flat_jax = util.split_list(args_and_out_cts_flat_jax,
                                                      [len(in_avals)])
    _, pullback_jax = jax.vjp(flattened_primal_fun_jax, *args_flat_jax)
    return pullback_jax(out_cts_flat_jax)

  vjp_in_avals = list(
      itertools.chain(in_avals,
                      map(lambda a: a.at_least_vspace(), out_avals)))

  # Expand in_shardings to all in_avals even not kept ones.
  all_in_shardings = [sharding_impls.UNSPECIFIED] * len(in_avals)
  for idx, in_s in zip(sorted(module_kept_var_idx),
                       in_shardings):  # type: ignore
    all_in_shardings[idx] = in_s  # type: ignore
  all_shardings = all_in_shardings + list(out_shardings)  # type: ignore
  # Cannot mix unspecified and specified shardings. Make the unspecified
  # ones replicated.
  specified_shardings = [
      s for s in all_shardings if not sharding_impls.is_unspecified(s)]

  vjp_in_shardings: Any  # The primal inputs followed by output cotangents
  vjp_out_shardings: Any  # The primal output cotangents
  if 0 == len(specified_shardings):
    vjp_in_shardings = sharding_impls.UNSPECIFIED
    vjp_out_shardings = sharding_impls.UNSPECIFIED
  else:
    if len(specified_shardings) < len(all_shardings):
      # There are some specified, but not all; pjit front-end does not liwk
      in_s = specified_shardings[0]  # pjit will enforce that all have same devices
      assert isinstance(in_s, sharding.XLACompatibleSharding)
      replicated_s = sharding.GSPMDSharding.get_replicated(in_s._device_assignment)
      all_shardings = [
          s if not sharding_impls.is_unspecified(s) else replicated_s
          for s in all_shardings]

    vjp_in_shardings = tuple(all_shardings)
    vjp_out_shardings = tuple(all_shardings[:len(in_avals)])
    if all(sharding_impls.is_unspecified(s) for s in vjp_out_shardings):
      vjp_out_shardings = sharding_impls.UNSPECIFIED

  if apply_jit:
    return pjit.pjit(fun_vjp_jax,
                     in_shardings=vjp_in_shardings,
                     out_shardings=vjp_out_shardings), vjp_in_avals
  else:
    assert vjp_in_shardings == sharding_impls.UNSPECIFIED
    assert vjp_out_shardings == sharding_impls.UNSPECIFIED
    return fun_vjp_jax, vjp_in_avals

def _export_native_vjp(primal_fun, primal: Exported) -> Exported:
  # Export the VJP of `primal_fun_jax`. See documentation for Exported.vjp
  fun_vjp_jax, vjp_in_avals = _get_vjp_fun(primal_fun,
                                           in_tree=primal.in_tree,
                                           module_kept_var_idx=primal.module_kept_var_idx,
                                           in_avals=primal.in_avals,
                                           in_shardings=primal.in_shardings,
                                           out_avals=primal.out_avals,
                                           out_shardings=primal.out_shardings,
                                           apply_jit=True)
  return export(fun_vjp_jax,
                lowering_platform=primal.lowering_platform,
                disabled_checks=primal.disabled_checks)(*vjp_in_avals)

### Importing

def call_exported(exported: Exported) -> Callable[..., jax.Array]:

  @jax.custom_vjp
  def f_flat(*args_flat):
    return call_exported_p.bind(*args_flat, exported=exported)

  def f_flat_vjp_fwd(*args_flat):
    # Return the primal arguments as the residual
    # TODO: keep as residuals only the arguments that are needed
    return f_flat(*args_flat), args_flat

  def f_flat_vjp_bwd(residual, ct_res_flat):
    args_flat = residual  # residual is the primal argument flat tuple
    exp_vjp = exported.vjp()
    in_ct_flat = call_exported(exp_vjp)(*args_flat, *ct_res_flat)
    return in_ct_flat

  f_flat.defvjp(f_flat_vjp_fwd, f_flat_vjp_bwd)

  def f_imported(*args, **kwargs):
    # since custom_vjp does not support kwargs, flatten the function first.
    args_flat, in_tree = tree_util.tree_flatten((args, kwargs))
    if in_tree != exported.in_tree:
      # Give errors with the precise tree difference; use fake leaves so we can
      # use tree_util.equality_errors.
      in_args = in_tree.unflatten([0] * in_tree.num_leaves)
      exp_in_args = exported.in_tree.unflatten([0] * exported.in_tree.num_leaves)

      msg = (
          "The invocation args and kwargs must have the same pytree structure "
          f"as when the function '{exported.fun_name}' was exported, but they "
          "have the following structural differences:\n" +
          ("\n".join(
             f"   - {shape_poly.args_kwargs_path_to_str(path)} is a {thing1} in the invocation and a "
             f"{thing2} when exported, so {explanation}.\n"
             for path, thing1, thing2, explanation
             in tree_util.equality_errors(in_args, exp_in_args))))
      raise ValueError(msg)

    res_flat = f_flat(*args_flat)
    return exported.out_tree.unflatten(res_flat)
  return f_imported


# A JAX primitive for invoking a serialized JAX function.
call_exported_p = core.Primitive("call_exported")
call_exported_p.multiple_results = True

@util.cache()
def _call_exported_abstract_eval(*in_avals: core.AbstractValue,
                                 exported: Exported) -> tuple[core.AbstractValue, ...]:
  exported_dim_vars = shape_poly.all_dim_vars(exported.in_avals)
  assert len(in_avals) == len(exported.in_avals)  # since the pytrees have the same structure
  # Check that the expected shapes match the actual ones
  for arg_idx, (exp_aval, actual_aval) in enumerate(zip(exported.in_avals, in_avals)):
    def pp_arg_dim(dim_idx: Optional[int]) -> str:
      return shape_poly.pretty_print_dimension_descriptor(exported.in_tree,
                                                          arg_idx, dim_idx)
    if len(exp_aval.shape) != len(actual_aval.shape):
      raise ValueError(
          f"Rank mismatch for {pp_arg_dim(None)}: expected {exp_aval.shape} "
          f"and called with {actual_aval.shape}")
    if exp_aval.dtype != actual_aval.dtype:
      raise ValueError(
          f"Dtype mismatch for {pp_arg_dim(None)}: expected {exp_aval.dtype} "
          f"and called with {actual_aval.dtype}")
    for dim_idx, aval_d in enumerate(exp_aval.shape):
      # If the exp_aval has a constant dimension then the actual argument must have
      # a matching constant dimension.
      if core.is_constant_dim(aval_d):
        if (not core.is_constant_dim(actual_aval.shape[dim_idx]) or
            aval_d != actual_aval.shape[dim_idx]):
          raise ValueError(
              f"Shape mismatch for {pp_arg_dim(dim_idx)} "
              "(expected same constant): "
              f"expected {exp_aval.shape} and called with {actual_aval.shape}")

  # Must express the exported_dim_vars in terms of the shapes in in_avals.
  solution, shape_constraints, synth_dim_vars = shape_poly.solve_dim_vars(
      exported.in_avals, args_kwargs_tree=exported.in_tree)
  synthetic_env = {vname: in_avals[arg_idx].shape[dim_idx]
                   for (vname, arg_idx, dim_idx) in synth_dim_vars}
  synthetic_eval = shape_poly.CachingShapeEvaluator(**synthetic_env)
  # We discharge all the constraints statically. This results in much simpler
  # composability (because we do not have to worry about the constraints of the
  # Exported called recursively; we only need to worry about entry-point
  # constraints). This also makes sense from a composibility point of view,
  # because we get the same errors if we invoke the exported module, or if we
  # trace the exported function. Consider for example, an exported module with
  # signature `f32[a, a] -> f32[a]`. If we invoke the module with an argument
  # `f32[c, d]` it is better to fail because `c == d` is inconclusive, than
  # succeed and add a compile-time check that `c == d`. In the latter case,
  # it would be ambiguous whether we should continue tracing with a result
  # a type `f32[c]` or `f32[d]`.
  shape_constraints.check_statically(synthetic_eval)
  exported_dim_values = [synthetic_eval.evaluate(solution[var])
                         for var in exported_dim_vars]
  return tuple(
      core.ShapedArray(core.evaluate_shape(out_aval.shape, exported_dim_vars,
                                           *exported_dim_values),
                       dtype=out_aval.dtype, weak_type=out_aval.weak_type,
                       named_shape=out_aval.named_shape)
      for out_aval in exported.out_avals)


call_exported_p.def_abstract_eval(_call_exported_abstract_eval)

def _call_exported_impl(*args, exported: Exported):
  return dispatch.apply_primitive(call_exported_p, *args, exported=exported)

call_exported_p.def_impl(_call_exported_impl)

def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
                            platform: str,
                            exported: Exported):
  # TODO: implement true multi-platform lowering for call_exported
  if (platform not in exported.lowering_platforms and
      DisabledSafetyCheck.platform() not in exported.disabled_checks):
    raise ValueError(
        f"The exported function '{exported.fun_name}' was lowered for "
        f"platforms '{exported.lowering_platforms}' but it is used "
        f"on '{platform}'.")

  if exported.uses_shape_polymorphism:
    ctx.module_context.shape_poly_state.uses_dim_vars = True

  submodule = ir.Module.parse(exported.mlir_module())
  symtab = ir.SymbolTable(submodule.operation)
  # The called function may have been exported with polymorphic shapes and called
  # now with more refined shapes. We insert hlo.ConvertOp to ensure the module
  # is valid.
  def convert_shape(x: ir.Value, x_aval: core.AbstractValue, new_aval: core.AbstractValue) -> ir.Value:
    new_ir_type = mlir.aval_to_ir_type(new_aval)
    if x.type != new_ir_type:
      return mlir.convert_hlo(ctx, x, x_aval, new_aval)
    else:
      return x

  callee_type = symtab["main"].type
  # TODO: maybe cache multiple calls
  fn = mlir.merge_mlir_modules(ctx.module_context.module,
                               f"call_exported_{exported.fun_name}",
                               submodule)
  kept_args = [
      convert_shape(a, a_aval, exported_in_aval)
      for i, (a, a_aval, exported_in_aval) in enumerate(zip(args, ctx.avals_in, exported.in_avals))
      if i in exported.module_kept_var_idx]
  if len(exported.lowering_platforms) > 1:
    # The exported module takes a platform index argument
    # TODO: implement proper handling of the platform_index when we are
    # in a multi-platform lowering context.
    platform_index = exported.lowering_platforms.index(platform)
    arg_width = callee_type.inputs[0].element_type.width
    assert arg_width in [32, 64]
    platform_index = np.int32(platform_index) if arg_width == 32 else np.int64(platform_index)  # type: ignore
    kept_args = [mlir.ir_constant(platform_index)] + kept_args
  call = func_dialect.CallOp(callee_type.results,
                             ir.FlatSymbolRefAttr.get(fn),
                             kept_args)
  # The ctx.avals_out already contain the abstract values refined by
  # _call_exported_abstract_eval.
  return tuple(
      convert_shape(out, out_aval, refined_out_aval)
      for out, out_aval, refined_out_aval in zip(call.results, exported.out_avals, ctx.avals_out))


for _p in ("cpu", "tpu", "cuda", "rocm"):
  mlir.register_lowering(call_exported_p,
                         functools.partial(_call_exported_lowering, platform=_p),
                         platform=_p)