code: def test_sharded_distribution_sharded_bijector(self):
      td = transformed_distribution.TransformedDistribution(
          normal.Normal(loc=0, scale=1), sigmoid.Sigmoid()
      )
      sharded_td = transformed_distribution.TransformedDistribution(
          sharded_lib.Sharded(
              normal.Normal(loc=0, scale=1),
              shard_axis_name=self.axis_name,
          ),
          sharded.Sharded(
              sigmoid.Sigmoid(), shard_axis_name=self.axis_name
          ),
      )
  
      x = self.evaluate(td.sample(test_lib.NUM_DEVICES, seed=self.key))
      sharded_x = self.shard_values(x)
  
      def true_lp_grad(x):
          def log_prob(x):
              return tf.reduce_sum(td.log_prob(x))
  
          lp, g = gradient.value_and_gradient(log_prob, (x,))
          return lp, g
  
      @tf.function
      def sharded_lp_grad(x):
          lp, g = gradient.value_and_gradient(sharded_td.log_prob, (x,))
          return lp, g
  
      true_lp, true_grad = self.evaluate(true_lp_grad(x))
      sharded_lp, sharded_grad = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(sharded_lp_grad, (sharded_x,))
          )
      )
      self.assertAllClose(true_lp, sharded_lp[0])
      self.assertAllClose(true_lp, sharded_lp[1])
      self.assertAllClose(true_grad, sharded_grad)
  
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 2 for calling the ShardedTest obj's underlying member method code with example usage: obj.test_sharded_distribution_sharded_bijector() and returns: lp,  g, lp,  g, tf