code: def test_correct_gradient_dtype_for_disconnected_variables(self):
      @tf.function(autograph=False)
      def run(x, y):
          def log_prob_parts(value):
              x, y = value
              return [
                  # These two RV's do not depend on each other.
                  normal.Normal(0.0, 1.0).log_prob(x),
                  normal.Normal(0.0, 1.0).log_prob(y),
              ]
  
          def log_prob(x, y):
              sharded_log_prob_parts = (
                  distribute_lib.make_sharded_log_prob_parts(
                      log_prob_parts, [None, self.axis_name]
                  )
              )
              parts = sharded_log_prob_parts([x, y])
              return tf.add_n(parts)
  
          return gradient.value_and_gradient(log_prob, (x, y))[1]
  
      sharded_x = self.shard_values(tf.range(4.0))
      sharded_y = sharded_x
      out_grads = self.per_replica_to_tensor(
          self.strategy_run(run, (sharded_x, sharded_y))
      )
      self.assertEqual(tf.float32, out_grads[0].dtype)
      self.assertEqual(tf.float32, out_grads[1].dtype)
  
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 13 for calling the LogProbPartsTest obj's underlying member method code with example usage: obj.test_correct_gradient_dtype_for_disconnected_variables() and returns: gradient, list, tf