code: def _custom_partitioning_infer_sharding_from_operands(
      arg_shapes, arg_shardings, result_shape, backend_string
  ):
      info = _sharding_callbacks[backend_string]
      if result_shape.is_tuple():
          result_shapes = result_shape.tuple_shapes()
      else:
          result_shapes = (result_shape,)
      result_sharding = info.infer_sharding_from_operands(
          *info.static_args,
          info.mesh,
          info.unflatten_arg_shapes(arg_shapes, arg_shardings),
          info.out_tree.unflatten(
              [_to_jax_shape(s) for s in result_shapes]
          ),
      )
      result_shardings = _flatten_sharding(
          info.out_tree, result_sharding, result_shapes
      )
      return _pack_result_sharding(result_shape, result_shardings)
  
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 9 for calling a global function with example usage: _custom_partitioning_infer_sharding_from_operands(arg_shapes, arg_shardings, result_shape, backend_string) and returns: _pack_result_sharding