code: def test_increment_log_prob(self):
      root = jdc.JointDistributionCoroutine.Root
      prior_mean = 3.0
      x_size = 100
  
      def custom_ll(w, x):
          return tf.reduce_sum(normal.Normal(w, 1.0).log_prob(x))
  
      def ulp_grad(w, x):
          @joint_distribution.JointDistributionCoroutine
          def sharded_model():
              w = yield root(normal.Normal(prior_mean, 1.0))
              yield root(
                  sharded.Sharded(
                      increment_log_prob.IncrementLogProb(
                          custom_ll(w, x)
                      ),
                      shard_axis_name=self.axis_name,
                  )
              )
  
          def ulp_fn(w):
              zeros = tf.zeros([x_size, 0])
              return sharded_model.unnormalized_log_prob(w, zeros)
  
          ulp, g = gradient.value_and_gradient(ulp_fn, (w,))
          return ulp, g
  
      def true_ulp_grad(w, x):
          @jdc.JointDistributionCoroutine
          def model():
              w = yield root(normal.Normal(prior_mean, 1.0))
              yield root(
                  increment_log_prob.IncrementLogProb(custom_ll(w, x))
              )
  
          def ulp_fn(w):
              zeros = tf.zeros([x_size, 0])
              return model.unnormalized_log_prob(w, zeros)
  
          ulp, g = gradient.value_and_gradient(ulp_fn, (w,))
          return ulp, g
  
      def test_w_x(w, x):
          sharded_x = self.shard_values(
              tf.reshape(x, [test_lib.NUM_DEVICES, -1])
          )
  
          lp, g = self.evaluate(
              self.per_replica_to_tensor(
                  self.strategy_run(
                      ulp_grad,
                      (
                          w,
                          sharded_x,
                      ),
                      in_axes=(None, 0),
                  )
              )
          )
          true_lp, true_g = self.evaluate(true_ulp_grad(w, x))
  
          self.assertAllClose(true_lp, lp[0])
          self.assertAllClose(true_g[0], g[0][0])
  
      w = tf.constant(4.0)
      zeros = tf.zeros([x_size])
      test_w_x(w, zeros)
      random_x = self.evaluate(
          normal.Normal(
              loc=tf.zeros([x_size]), scale=tf.ones([x_size])
          ).sample(seed=self.key)
      )
      test_w_x(w, random_x)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.distribute import joint_distribution
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.experimental.distributions import increment_log_prob
  from tensorflow_probability.python.internal import distribute_test_lib
  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 ShardTest obj's underlying member method code with example usage: obj.test_increment_log_prob() and returns: tf, ulp,  g, ulp,  g, sharded_model, model