code: def make_sharded_log_prob_parts(log_prob_parts_fn, axis_names):
      """Constructs a log prob parts function that all-reduces over terms.
  
      Given a log_prob_parts function, this function will return a new one that
      includes all-reduce sums over terms according to the `is_sharded` property. It
      will also add all-reduce sums for the gradient of sharded terms w.r.t.
      unsharded terms.
  
      Args:
        log_prob_parts_fn: a callable that takes in a structured value and returns a
          structure of log densities for each of the terms, that when summed returns
          a locally correct log-density.
        axis_names: a structure of values that matches the input and output of
          `log_prob_parts_fn`. Each value in `axis_names` is either `None, a string
          name of a mapped axis in the JAX backend or any non-`None` value in TF
          backend, or an iterable thereof corresponding to multiple sharding axes.
          If the `axis_name` is not `None`, the returned function will add
          all-reduce sum(s) for its term in the log prob calculation. If it is
          `None`, the returned function will have an all-reduce sum over the
          gradient of sharded terms w.r.t. to the unsharded value.
  
      Returns:
        A new log prob parts function that can be run inside of a strategy.
      """
  
      def sharded_log_prob_parts_fn(*args):
          out_dtype = tf.nest.map_structure(lambda x: x.dtype, args)
          if len(args) == 1:
              out_dtype = out_dtype[0]
          return make_psum_function(
              log_prob_parts_fn, (axis_names,), axis_names, out_dtype
          )(*args)
  
      return sharded_log_prob_parts_fn
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow.python.util import nest
label_desc: Use image node_id 20 for calling a global function with example usage: make_sharded_log_prob_parts(log_prob_parts_fn, axis_names) and returns: sharded_log_prob_parts_fn, make_psum_function