code: 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 not re.match(
                      check_sharding_pattern,
                      ir.StringAttr(sharding).value,
                  ):
                      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":
              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)
  
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 10 for calling a global function with example usage: _check_module(mod) without return types