code: def _custom_partitioning_partition(
      arg_shapes,
      arg_shardings,
      result_shape,
      result_sharding,
      backend_string,
  ):
      info = _sharding_callbacks[backend_string]
      if result_shape.is_tuple():
          result_shapes = result_shape.tuple_shapes()
          result_shardings = result_sharding.tuple_elements()
      else:
          result_shapes = (result_shape,)
          result_shardings = (result_sharding,)
      mesh, lower_fn, result_sharding, arg_shardings = info.partition(
          *info.static_args,
          info.mesh,
          info.unflatten_arg_shapes(arg_shapes, arg_shardings),
          info.out_tree.unflatten(
              [
                  info.unflatten_arg_shape(s, sharding)
                  for s, sharding in zip(
                      result_shapes, result_shardings
                  )
              ]
          ),
      )
      module_context = info.module_context
  
      result_shardings = _flatten_sharding(
          info.out_tree, result_sharding, result_shapes
      )
      arg_shardings = _flatten_sharding(
          info.in_tree, arg_shardings, arg_shapes
      )
      tiled_args = [
          _to_jax_shape(sharding.tile(s))
          for sharding, s in zip(arg_shardings, arg_shapes)
      ]
      tiled_results = [
          _to_jax_shape(sharding.tile(s))
          for sharding, s in zip(result_shardings, result_shapes)
      ]
      closed_jaxpr = jax.make_jaxpr(
          lower_fn, axis_env=list(mesh.shape.items())
      )(*tiled_args)
      if closed_jaxpr.out_avals != tiled_results:
          raise ValueError(
              "Mismatch in result shapes. %s vs %s"
              % (repr(closed_jaxpr.out_avals), repr(tiled_results))
          )
      axis_context = sharding_impls.SPMDAxisContext(mesh)
      built = mlir.build_xla_computation_helper(
          closed_jaxpr,
          name="tmp_xla_computation",
          platforms=module_context.platforms,
          backend_or_name=module_context.backend_or_name,
          axis_context=axis_context.extend_manual(
              frozenset(mesh.axis_names)
          ),
      )
      result_sharding = _pack_result_sharding(
          result_shape, result_shardings
      )
      return built, arg_shardings, result_sharding
  
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 8 for calling a global function with example usage: _custom_partitioning_partition(arg_shapes, arg_shardings, result_shape, result_sharding, backend_string) and returns: built,  arg_shardings,  result_sharding