code: def test_correct_log_prob_for_local_variable(self):
      def run(x, data):
          def log_prob_parts(value):
              x, data = value
              return [
                  normal.Normal(0.0, 1.0).log_prob(x),
                  tf.reduce_sum(normal.Normal(x, 1.0).log_prob(data)),
              ]
  
          sharded_log_prob_parts = (
              distribute_lib.make_sharded_log_prob_parts(
                  log_prob_parts, [self.axis_name, self.axis_name]
              )
          )
  
          return sharded_log_prob_parts([x, data])
  
      x = tf.zeros(4)
      sharded_x = self.shard_values(x)
      data = tf.ones(4)
      sharded_data = self.shard_values(data)
      out_parts = self.per_replica_to_tensor(
          self.strategy_run(run, (sharded_x, sharded_data))
      )
  
      self.assertAllEqualNested(
          self.evaluate(out_parts),
          self.evaluate(
              [
                  tf.ones(4)
                  * tf.reduce_sum(normal.Normal(0.0, 1.0).log_prob(x)),
                  tf.ones(4)
                  * tf.reduce_sum(
                      normal.Normal(0.0, 1.0).log_prob(data)
                  ),
              ]
          ),
      )
  
imports: import functools
  import itertools
  from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 5 for calling the LogProbPartsTest obj's underlying member method code with example usage: obj.test_correct_log_prob_for_local_variable() and returns: sharded_log_prob_parts, list