code: def test_sharded_log_det_jacobian(self):
      def log_prob(y):
          return tf.reduce_sum(beta.Beta(2.0, 2.0).log_prob(y))
  
      def transform_log_prob(log_prob, bijector):
          def new_log_prob(x):
              y = bijector.forward(x)
              return log_prob(y) + bijector.forward_log_det_jacobian(
                  x, len(x.shape)
              )
  
          return new_log_prob
  
      @tf.function
      def lp_grad(x):
          untransformed_log_prob = (
              distribute_lib.make_sharded_log_prob_parts(
                  log_prob, self.axis_name
              )
          )
          transformed_log_prob = transform_log_prob(
              untransformed_log_prob,
              sharded.Sharded(
                  sigmoid.Sigmoid(), shard_axis_name=self.axis_name
              ),
          )
          lp, g = gradient.value_and_gradient(
              transformed_log_prob, (x,)
          )
          return lp, g
  
      def true_lp_grad(x):
          transformed_log_prob = transform_log_prob(
              log_prob, sigmoid.Sigmoid()
          )
          lp, g = gradient.value_and_gradient(
              transformed_log_prob, (x,)
          )
          return lp, g
  
      y = tf.convert_to_tensor(
          np.linspace(0.2, 0.7, test_lib.NUM_DEVICES, dtype=np.float32)
      )
      x = self.evaluate(sigmoid.Sigmoid().inverse(y))
      sharded_x = self.shard_values(x)
      lp, g = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(lp_grad, (sharded_x,))
          )
      )
      true_lp, true_g = self.evaluate(true_lp_grad(x))
      self.assertAllClose(true_lp, lp[0])
      self.assertAllClose(true_lp, lp[1])
      self.assertAllClose(true_g, g)
  
file_path: probability/tensorflow_probability/python/experimental/bijectors/sharded_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.bijectors import sharded
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.internal import distribute_lib
  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 1 for calling the ShardedTest obj's underlying member method code with example usage: obj.test_sharded_log_det_jacobian() and returns: tf, new_log_prob, lp,  g, lp,  g, unknown