code: def _implement_sharded_lp_fn(fn_name):
      """Implements log_prob or unnormalized_log_prob."""
  
      def lp_fn(self, x):
          lp = getattr(self.distribution, fn_name)(x)
          return distribute_lib.psum(
              lp, self.experimental_shard_axis_names
          )
  
      lp_fn.__name__ = f"_{fn_name}"
      return lp_fn
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.experimental.bijectors import sharded
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 1 for calling a global function with example usage: _implement_sharded_lp_fn(fn_name) and returns: lp_fn, distribute_lib