code: def _custom_partitioning_lowering_rule(
      ctx: mlir.LoweringRuleContext,
      *values,
      call,
      in_tree,
      out_tree,
      propagate_user_sharding,
      partition,
      infer_sharding_from_operands,
      decode_shardings,
      static_args,
  ):
      mesh = mesh_lib.thread_resources.env.physical_mesh
      axis_context = ctx.module_context.axis_context
      if isinstance(
          axis_context, sharding_impls.SPMDAxisContext
      ) and set(axis_context.manual_axes) == set(
          axis_context.mesh.axis_names
      ):
          return mlir.lower_fun(
              core.jaxpr_as_fun(call), multiple_results=True
          )(ctx, *values)
  
      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:
          devices = None
  
      if not devices or len(devices) == 1:
          return mlir.lower_fun(
              core.jaxpr_as_fun(call), multiple_results=True
          )(ctx, *values)
  
      def to_mesh_pspec_sharding(
          hlo_sharding: xc.HloSharding | None, ndim
      ):
          if hlo_sharding is None:
              return hlo_sharding
          if mesh.empty or not decode_shardings:
              assert devices is not None
              return _op_sharding_to_pos_sharding(hlo_sharding, devices)
          pspec = sharding_impls.parse_flatten_op_sharding(
              hlo_sharding, mesh
          )[0].get_partition_spec()
          pspec = jax.sharding.PartitionSpec(
              *pspec, *((None,) * (ndim - len(pspec)))
          )
          return jax.sharding.NamedSharding(mesh, pspec)
  
      sharding_callback_info = _ShardingCallbackInfo(
          propagate_user_sharding,
          partition,
          to_mesh_pspec_sharding,
          in_tree,
          out_tree,
          infer_sharding_from_operands,
          ctx.module_context,
          mesh,
          static_args,
      )
      key = str(id(sharding_callback_info))
      # TODO(parkers): Remove bytes registration when xla_extension_version > 211
      _sharding_callbacks[key] = sharding_callback_info
      _sharding_callbacks[bytes(key, "utf8")] = sharding_callback_info
      # We need to make sure `sharding_callback_info` is still alive when the SPMD
      # partitioner runs so we keep it alive by attaching it to the executable.
      ctx.module_context.add_keepalive(sharding_callback_info)
  
      result_types = [mlir.aval_to_ir_type(s) for s in call.out_avals]
      out = hlo.CustomCallOp(
          result_types,
          list(values),
          call_target_name=ir.StringAttr.get(
              _CUSTOM_PARTITIONING_CALL_NAME
          ),
          has_side_effect=ir.BoolAttr.get(False),
          api_version=mlir.i32_attr(2),
          called_computations=ir.ArrayAttr.get([]),
          backend_config=ir.StringAttr.get(key),
          operand_layouts=None,
          result_layouts=None,
      )
      return out.results
  
imports: from __future__ import annotations
  import inspect
  from typing import Optional
  import weakref
  import jax
  from jax._src import core
  from jax import tree_util
  from jax._src import linear_util
  from jax._src import sharding_impls
  from jax.errors import UnexpectedTracerError
  from jax._src import mesh
  from jax._src import dispatch
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.lib.mlir import ir
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.sharding_impls import _op_sharding_to_pos_sharding
  from jax._src import custom_api_util
  from jax._src.lib import xla_client
  from jax._src.api_util import flatten_fun_nokwargs, argnums_partial
label_desc: Use image node_id 13 for calling a global function with example usage: _custom_partitioning_lowering_rule(ctx) and returns: out, mlir, mlir, jax, hlo_sharding, _op_sharding_to_pos_sharding