code: def test_correct_gradient_for_local_integer_variable(self):
      @tf.function(autograph=False)
      def run(x, data):
          def log_prob_parts(value):
              x, data = value
              return [
                  normal.Normal(0.0, 1.0).log_prob(x),
                  bernoulli.Bernoulli(logits=x).log_prob(data),
              ]
  
          def log_prob(x):
              sharded_log_prob_parts = (
                  distribute_lib.make_sharded_log_prob_parts(
                      log_prob_parts, [self.axis_name, self.axis_name]
                  )
              )
              parts = sharded_log_prob_parts([x, data])
              return tf.add_n(parts)
  
          return gradient.value_and_gradient(log_prob, x)[1]
  
      x = tf.range(4.0)
      sharded_x = self.shard_values(x)
      data = tf.ones(4, tf.int32)
      sharded_data = self.shard_values(data)
      out_grads = self.per_replica_to_tensor(
          self.strategy_run(run, (sharded_x, sharded_data))
      )
  
      def true_log_prob(x):
          return tf.reduce_sum(
              normal.Normal(0.0, 1.0).log_prob(x)
          ) + tf.reduce_sum(
              bernoulli.Bernoulli(logits=x).log_prob(data)
          )
  
      true_grad = self.evaluate(
          gradient.value_and_gradient(true_log_prob, x)[1]
      )
  
      self.assertAllEqualNested(self.evaluate(out_grads), true_grad)
  
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 12 for calling the LogProbPartsTest obj's underlying member method code with example usage: obj.test_correct_gradient_for_local_integer_variable() and returns: gradient, unknown, list, tf