code: 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",
          "in_layouts",
          "out_layouts",
          "all_args_info",
      ]
      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", 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)
          )
  
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 9 for calling a global function with example usage: _check_lowering(lowering) without return types