code: def export(
      fun_jax: Callable,
      *,
      lowering_platforms: Sequence[str] | None = None,
      disabled_checks: Sequence[DisabledSafetyCheck] = (),
  ) -> Callable[..., Exported]:
      """Exports native serialization for a JAX function.
  
      Args:
        fun_jax: the function to lower and serialize.
        lowering_platforms:
            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.value
      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
  
          if lowering_platforms is not None:
              actual_lowering_platforms = tuple(lowering_platforms)
          else:
              actual_lowering_platforms = (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 >= _VERSION_START_SUPPORT_SHAPE_ASSERTIONS
          )  # 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
              )
              replace_tokens_with_dummy = not _keep_main_tokens(version)
              lowered = wrapped_fun_jax.lower(
                  *args_specs,
                  **kwargs_specs,
                  _experimental_lowering_parameters=mlir.LoweringParameters(
                      platforms=actual_lowering_platforms,
                      replace_tokens_with_dummy=replace_tokens_with_dummy,
                  ),
              )
  
              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,
                      serialization_version=version,
                  )
          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 = _module_to_bytecode(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):
              logmsg = (
                  f"version={version} "
                  f"lowering_platforms={actual_lowering_platforms} "
                  f"disabled_checks={disabled_checks}"
              )
              logging.info("Lowered JAX module: %s\n", logmsg)
              if dumped_to := mlir.dump_module_to_file(
                  mlir_module, "export"
              ):
                  logging.info(
                      "Dumped the exported MLIR module to %s", dumped_to
                  )
  
          _check_module(
              mlir_module,
              allow_non_replicated_sharding=allow_non_replicated_sharding,
              disabled_checks=disabled_checks,
          )
  
          ordered_effects = tuple(
              lowering.compile_args["ordered_effects"]
          )
          unordered_effects = tuple(
              lowering.compile_args["unordered_effects"]
          )
          if version < _VERSION_START_SUPPORT_EFFECTS_WITH_REAL_TOKENS:
              ordered_effects = unordered_effects = ()
  
          nr_devices = len(lowering.compile_args["device_assignment"])
  
          def export_sharding(
              s: LoweringSharding, aval: core.ShapedArray
          ) -> Sharding:
              if sharding_impls.is_unspecified(s):
                  return None
              return s._to_xla_hlo_sharding(aval.ndim)  # type: ignore[union-attr]
  
          all_in_shardings = expand_in_shardings(
              lowering.compile_args["in_shardings"],
              module_kept_var_idx,
              len(args_avals_flat),
          )
          in_shardings = tuple(
              export_sharding(s, aval)
              for s, aval in zip(all_in_shardings, args_avals_flat)
          )
          out_shardings = tuple(
              export_sharding(s, aval)
              for s, aval in zip(
                  lowering.compile_args["out_shardings"], out_avals_flat
              )
          )
          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=in_shardings,
              out_shardings=out_shardings,
              nr_devices=nr_devices,
              lowering_platforms=actual_lowering_platforms,
              ordered_effects=ordered_effects,
              unordered_effects=unordered_effects,
              disabled_safety_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,
              mlir_module_serialization_version=version,  # type: ignore
              _get_vjp=lambda exported: _export_native_vjp(
                  fun_jax, exported
              ),
          )
  
      return do_export
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import copy
  import dataclasses
  import functools
  import itertools
  import re
  from typing import Any, Callable, Union
  import warnings
  from absl import logging
  import numpy
  import jax
  from jax import sharding
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  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
  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
  from jax.experimental.export import _shape_poly
label_desc: Use image node_id 6 for calling a global function with example usage: export(fun_jax) and returns: do_export, Exported, s, None