code: def test_correct_gradient_for_global_and_local_variable_dict(self):
      @tf.function(autograph=False)
      def run(w, x, data):
          def log_prob_parts(value):
              return {
                  "w": normal.Normal(0.0, 1.0).log_prob(value["w"]),
                  "x": normal.Normal(value["w"], 1.0).log_prob(
                      value["x"]
                  ),
                  "data": normal.Normal(value["x"], 1.0).log_prob(
                      value["data"]
                  ),
              }
  
          def log_prob(*value):
              w, x = value
              sharded_log_prob_parts = (
                  distribute_lib.make_sharded_log_prob_parts(
                      log_prob_parts,
                      {
                          "w": None,
                          "x": self.axis_name,
                          "data": self.axis_name,
                      },
                  )
              )
              parts = sharded_log_prob_parts(
                  {"w": w, "x": x, "data": data}
              )
              return tf.add_n(tf.nest.flatten(parts))
  
          return gradient.value_and_gradient(log_prob, [w, x])[1]
  
      w = tf.constant(1.0)
      x = tf.range(4.0)
      sharded_x = self.shard_values(x)
      data = 2 * tf.ones(4)
      sharded_data = self.shard_values(data)
      out_grads = self.per_replica_to_tensor(
          self.strategy_run(
              run, (w, sharded_x, sharded_data), in_axes=(None, 0, 0)
          )
      )
  
      def true_log_prob(*value):
          w, x = value
          return (
              normal.Normal(0.0, 1.0).log_prob(w)
              + tf.reduce_sum(normal.Normal(w, 1.0).log_prob(x))
              + tf.reduce_sum(normal.Normal(x, 1.0).log_prob(data))
          )
  
      true_grad = gradient.value_and_gradient(true_log_prob, [w, x])[1]
      true_grad[0] = tf.ones(4) * true_grad[0]
  
      self.assertAllEqualNested(
          self.evaluate(out_grads), self.evaluate(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 11 for calling the LogProbPartsTest obj's underlying member method code with example usage: obj.test_correct_gradient_for_global_and_local_variable_dict() and returns: gradient, unknown, dict, tf