code: def _custom_partitioning_propagate_user_sharding(
      user_sharding, shape, backend_string
  ):
      info = _sharding_callbacks[backend_string]
      if info.propagate_user_sharding is None:
          return user_sharding
      if shape.is_tuple():
          user_shapes = shape.tuple_shapes()
          user_shardings = user_sharding.tuple_elements()
      else:
          user_shapes = (shape,)
          user_shardings = (user_sharding,)
      user_shape = info.out_tree.unflatten(
          [
              info.unflatten_arg_shape(s, sharding)
              for s, sharding in zip(user_shapes, user_shardings)
          ]
      )
      result_sharding = info.propagate_user_sharding(
          *info.static_args, info.mesh, user_shape
      )
      result_shardings = _flatten_sharding(
          info.out_tree, result_sharding, user_shapes
      )
      return _pack_result_sharding(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 6 for calling a global function with example usage: _custom_partitioning_propagate_user_sharding(user_sharding, shape, backend_string) and returns: _pack_result_sharding, user_sharding