code: def _inspect_sharding_lowering_rule(
      ctx: mlir.LoweringRuleContext, value, *, callback
  ):
      mesh = mesh_lib.thread_resources.env.physical_mesh
      axis_context = ctx.module_context.axis_context
  
      if isinstance(axis_context, sharding_impls.ShardingContext):
          devices = axis_context.device_assignment
          if devices is None:
              raise AssertionError(
                  "Please file a bug at https://github.com/google/jax/issues"
              )
      elif isinstance(axis_context, sharding_impls.SPMDAxisContext):
          devices = axis_context.mesh._flat_devices_tuple
      else:
          raise NotImplementedError(type(axis_context))
      assert devices is not None
  
      # If we have a nontrivial parallel computation, we need to wait until the SPMD
      # partitioner calls back with the `HloSharding.
      def _hlo_sharding_callback(hlo_sharding: xc.HloSharding):
          if mesh.empty:
              return callback(
                  sharding_impls._op_sharding_to_pos_sharding(
                      hlo_sharding, devices
                  )
              )
          pspec = parse_flatten_op_sharding(hlo_sharding, mesh)[
              0
          ].get_partition_spec()
          return callback(NamedSharding(mesh, pspec))
  
      if len(devices) == 1:
          # If we only have one device in our computation, we can construct a
          # replicated HloSharding and call it right now.
          _hlo_sharding_callback(
              sharding_impls.get_replicated_hlo_sharding()
          )
          return []
  
      key = xc.encode_inspect_sharding_callback(_hlo_sharding_callback)
      # We need to make sure `_hlo_sharding_callback` is still alive when the SPMD
      # partitioner runs so we keep it alive by attaching it to the executable.    #
      ctx.module_context.add_keepalive(_hlo_sharding_callback)
  
      hlo.CustomCallOp(
          [value.type],
          [value],
          call_target_name=ir.StringAttr.get(
              _INSPECT_SHARDING_CALL_NAME
          ),
          has_side_effect=ir.BoolAttr.get(True),
          api_version=mlir.i32_attr(1),
          called_computations=ir.ArrayAttr.get([]),
          backend_config=ir.StringAttr.get(key),
          operand_layouts=None,
          result_layouts=None,
      )
      return []
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import string
  import sys
  from typing import Any, Callable, Union
  import weakref
  import numpy
  import jax.numpy
  from jax import lax
  from jax._src import core
  from jax._src import effects
  from jax._src import mesh
  from jax._src import sharding_impls
  from jax._src import dispatch
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  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.sharding import Sharding
  from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
label_desc: Use image node_id 15 for calling a global function with example usage: _inspect_sharding_lowering_rule(ctx, value) and returns: list, callback, list, callback